Skip to content

Commit 077726f

Browse files
committed
updated
1 parent 4abfce2 commit 077726f

File tree

1 file changed

+105
-20
lines changed

1 file changed

+105
-20
lines changed

data_structures/arrays/median_two_array.py

+105-20
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,18 @@
1-
"""
2-
https://www.enjoyalgorithms.com/blog/median-of-two-sorted-arrays
3-
"""
4-
5-
61
def find_median_sorted_arrays(nums1: list[int], nums2: list[int]) -> float:
72
"""
8-
Find the median of two arrays.
3+
Find the median of two sorted arrays.
4+
5+
This implementation uses binary search to achieve an optimal solution in O(log(min(m, n))) time.
96
107
Args:
11-
nums1: The first array.
12-
nums2: The second array.
8+
nums1: The first sorted array.
9+
nums2: The second sorted array.
1310
1411
Returns:
15-
The median of the two arrays.
12+
The median of the two sorted arrays.
13+
14+
Raises:
15+
ValueError: If both input arrays are empty.
1616
1717
Examples:
1818
>>> find_median_sorted_arrays([1, 3], [2])
@@ -25,7 +25,7 @@ def find_median_sorted_arrays(nums1: list[int], nums2: list[int]) -> float:
2525
0.0
2626
2727
>>> find_median_sorted_arrays([], [])
28-
Traceback (most recent call last):
28+
Traceback (most recent last):
2929
...
3030
ValueError: Both input arrays are empty.
3131
@@ -41,19 +41,104 @@ def find_median_sorted_arrays(nums1: list[int], nums2: list[int]) -> float:
4141
if not nums1 and not nums2:
4242
raise ValueError("Both input arrays are empty.")
4343

44-
# Merge the arrays into a single sorted array.
45-
merged = sorted(nums1 + nums2)
46-
total = len(merged)
44+
# Insert 1: Ensure the smaller array is nums1
45+
if len(nums1) > len(nums2):
46+
nums1, nums2 = nums2, nums1
47+
48+
m, n = len(nums1), len(nums2)
49+
low, high = 0, m
50+
51+
while low <= high:
52+
partition1 = (low + high) // 2
53+
partition2 = (m + n + 1) // 2 - partition1
54+
55+
# Handle edges of the arrays with infinities
56+
max_left1 = float('-inf') if partition1 == 0 else nums1[partition1 - 1]
57+
min_right1 = float('inf') if partition1 == m else nums1[partition1]
4758

48-
if total % 2 == 1: # If the total number of elements is odd
49-
return float(merged[total // 2]) # then return the middle element
59+
max_left2 = float('-inf') if partition2 == 0 else nums2[partition2 - 1]
60+
min_right2 = float('inf') if partition2 == n else nums2[partition2]
5061

51-
# If the total number of elements is even, calculate
52-
# the average of the two middle elements as the median.
53-
middle1 = merged[total // 2 - 1]
54-
middle2 = merged[total // 2]
55-
return (float(middle1) + float(middle2)) / 2.0
62+
# Insert 2: Debugging: Log partition indices and elements (useful for large arrays)
63+
# print(f"partition1: {partition1}, partition2: {partition2}")
64+
# print(f"max_left1: {max_left1}, min_right1: {min_right1}")
65+
# print(f"max_left2: {max_left2}, min_right2: {min_right2}")
5666

67+
# Check if we have found the correct partition
68+
if max_left1 <= min_right2 and max_left2 <= min_right1:
69+
if (m + n) % 2 == 1:
70+
return max(max_left1, max_left2)
71+
else:
72+
return (max(max_left1, max_left2) + min(min_right1, min_right2)) / 2.0
73+
elif max_left1 > min_right2:
74+
high = partition1 - 1
75+
else:
76+
low = partition1 + 1
77+
78+
# Insert 3: Remove redundant exception, already handled at the beginning
79+
# raise ValueError("Input arrays are not sorted.") # This line is no longer necessary.
80+
81+
def merge_sorted_arrays(nums1: list[int], nums2: list[int]) -> list[int]:
82+
"""
83+
Merge two sorted arrays into a single sorted array in O(m + n) time.
84+
85+
Args:
86+
nums1: The first sorted array.
87+
nums2: The second sorted array.
88+
89+
Returns:
90+
A new sorted array containing all elements of nums1 and nums2.
91+
92+
Examples:
93+
>>> merge_sorted_arrays([1, 3], [2])
94+
[1, 2, 3]
95+
96+
>>> merge_sorted_arrays([1, 2], [3, 4])
97+
[1, 2, 3, 4]
98+
99+
>>> merge_sorted_arrays([0, 0], [0, 0])
100+
[0, 0, 0, 0]
101+
102+
>>> merge_sorted_arrays([], [1])
103+
[1]
104+
105+
>>> merge_sorted_arrays([], [])
106+
[]
107+
"""
108+
# Insert 4: Edge case: If one array is empty, just return the other array
109+
if not nums1:
110+
return nums2
111+
if not nums2:
112+
return nums1
113+
114+
i, j = 0, 0
115+
merged = []
116+
117+
while i < len(nums1) and j < len(nums2):
118+
if nums1[i] < nums2[j]:
119+
merged.append(nums1[i])
120+
i += 1
121+
else:
122+
merged.append(nums2[j])
123+
j += 1
124+
125+
# Insert 5: Append remaining elements from either nums1 or nums2
126+
merged.extend(nums1[i:])
127+
merged.extend(nums2[j:])
128+
return merged
129+
130+
131+
def is_sorted(nums: list[int]) -> bool:
132+
"""
133+
Helper function to check if the array is sorted.
134+
135+
Args:
136+
nums: The array to check.
137+
138+
Returns:
139+
True if the array is sorted, False otherwise.
140+
"""
141+
return all(nums[i] <= nums[i + 1] for i in range(len(nums) - 1))
57142

58143
if __name__ == "__main__":
59144
import doctest

0 commit comments

Comments
 (0)