Add median of two sorted arrays implementation #13717

divide_and_conquer/median_of_two_sorted_arrays.py

@@ -0,0 +1,88 @@
+"""
+Median of Two Sorted Arrays
+
+Given two sorted arrays nums1 and nums2 of size m and n respectively, return the median of the two sorted arrays.
+
+The overall run time complexity should be O(log(m + n)).
+
+Examples:
+ Example 1:
+ Input: nums1 = [1, 3], nums2 = [2]
+ Output: 2.00000
+ Explanation: merged array = [1, 2, 3] and median is 2.
+
+ Example 2:
+ Input: nums1 = [1, 2], nums2 = [3, 4]
+ Output: 2.50000
+ Explanation: merged array = [1, 2, 3, 4] and median is (2 + 3) / 2 = 2.5.
+
+Constraints:
+ * nums1.length == m
+ * nums2.length == n
+ * 0 <= m <= 1000
+ * 0 <= n <= 1000
+ * 1 <= m + n <= 2000
+ * -10^6 <= nums1[i], nums2[i] <= 10^6
+
+Implementation: Divide and Conquer (Binary Search on partitions).
+Time Complexity: O(log(min(m, n)))
+Space Complexity: O(1)
+"""
+
+from typing import List
+
+
+def find_median_sorted_arrays(nums1: List[int], nums2: List[int]) -> float:
+ """
+ Find the median of two sorted arrays using binary search on partitions.
+
+ Args:
+ nums1 (List[int]): First sorted array
+ nums2 (List[int]): Second sorted array
+
+ Returns:
+ float: Median of the two arrays
+ """
+ # Ensure nums1 is the smaller array
+ if len(nums1) > len(nums2):
+ nums1, nums2 = nums2, nums1
+
+ m, n = len(nums1), len(nums2)
+ left, right = 0, m
+
+ # Binary search for the partition in nums1
+ while left <= right:
+ i = (left + right) // 2 # Partition in nums1
+ j = (m + n + 1) // 2 - i # Partition in nums2
+
+ # Edge cases for partitions
+ left1 = float('-inf') if i == 0 else nums1[i - 1]
+ right1 = float('inf') if i == m else nums1[i]
+ left2 = float('-inf') if j == 0 else nums2[j - 1]
+ right2 = float('inf') if j == n else nums2[j]
+
+ # Check if this partition is correct
+ if left1 <= right2 and left2 <= right1:
+ # Correct partition found
+ if (m + n) % 2 == 0:
+ # Even length: average of max(lefts) and min(rights)
+ return (max(left1, left2) + min(right1, right2)) / 2
+ else:
+ # Odd length: max of lefts
+ return max(left1, left2)
+ elif left1 > right2:
+ # Move partition left in nums1
+ right = i - 1
+ else:
+ # Move partition right in nums1
+ left = i + 1
+
+ # Should not reach here if inputs are valid
+ raise ValueError("Input arrays are not sorted or invalid")
+
+
+# Optional: Add simple tests
+if __name__ == "__main__":
+ assert abs(find_median_sorted_arrays([1, 3], [2]) - 2.0) < 1e-5
+ assert abs(find_median_sorted_arrays([1, 2], [3, 4]) - 2.5) < 1e-5
+ print("All tests passed!")