LeetCode (Algorithms#4) ‒ Median of Two Sorted Arrays

Problem

There are two sorted arrays nums1 and nums2 of size m and n respectively.

Find the median of the two sorted arrays. The overall run time complexity should be O(log(m+n)).

Example 1:

nums1 = [1, 3]
nums2 = [2]

The median is 2.0

Example 2:

nums1 = [1, 2]
nums2 = [3, 4]

The median is (2 + 3)/2 = 2.5

Solution

The idea behind this problem is quite simple. To find the median in O(log(n+m)), we need to perform some kind of binary search concurrently on both arrays. The difficulty is in making sure to consider all possible cases. Overall, there is a general case that performs the binary search, and a base case to chose the median among four or less remaining elements. The code below has the different cases commented. The idea is to keep track of how many elements have been removed from each side, and repeatedly split either array in half such that we remove the smaller or larger elements from the left or right side, respectively.

This is certainly not the most clean or efficient solution to this problem as there are many cases that need to be explicitly handled. Another solution is to use a recursive divide and conquer approach where you determine the two array's medians, and then recurse on the halves of the arrays that bring the medians closer (or equal to each other).

Time complexity: O(n+m)
Space complexity: O(n^2)

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86

public double findMedianSortedArrays(int[] nums1, int[] nums2) { int len1 = nums1.length, len2 = nums2.length, len = len1 + len2; boolean evenlen = (len1 + len2) % 2 == 0; int l1 = 0, l2 = 0, r1 = len1 - 1, r2 = len2 - 1, piv1, piv2; int leftCt, rightCt; if (len1 == 0 && len2 == 0) { return 0; } else if (len1 == 0) { return (evenlen ? (nums2[r2/2] + nums2[r2/2+1]) / 2. : nums2[r2/2]); } else if (len2 == 0) { return (evenlen ? (nums1[r1/2] + nums1[r1/2+1]) / 2. : nums1[r1/2]); } // Reduce subsets until 1 or 2 medians remain while (true) { piv1 = l1 + (r1 - l1) / 2; piv2 = l2 + (r2 - l2) / 2; boolean nonZero1 = r1 - l1 >= 0, nonZero2 = r2 - l2 >= 0; int sz = (r1 - l1) + (r2 - l2); if (sz <= 2) { while (true) { leftCt = l1 + l2; rightCt = (len1 - 1 - r1) + (len2 - 1 - r2); if (evenlen && (r1 - l1) + (r2 - l2) == 0) { if (r1 - l1 == -1) return (nums2[l2] + nums2[r2]) / 2.; else if (r2 - l2 == -1) return (nums1[l1] + nums1[r1]) / 2.; else return (nums1[l1] + nums2[l2]) / 2.; } else if (!evenlen && (r1 - l1) + (r2 - l2) == -1) { if (r1 - l1 == -1) return nums2[l2]; else return nums1[l1]; } if (leftCt < rightCt) { // remove a smaller value if (r1 - l1 >= 0 && (r2 - l2 < 0 || nums1[l1] < nums2[l2])) l1++; else l2++; } else { // remove a larger value if (r1 - l1 >= 0 && (r2 - l2 < 0 || nums1[r1] > nums2[r2])) r1--; else r2--; } } } else { // General case leftCt = l1 + l2; rightCt = (len1 - 1 - r1) + (len2 - 1 - r2); if (leftCt > rightCt) { // remove from right // reduce array with max value on right if (!nonZero1 || !nonZero2) { // One array is empty if (nonZero1) r1 = piv1; else r2 = piv2; } else { if (piv1 < r1 && piv2 < r2) { // Both arrays have elements right of pivots if (nums1[piv1+1] > nums2[piv2+1]) r1 = piv1; else r2 = piv2; } else { if (piv1 == l1) { if (nums1[piv1] > nums2[piv2+1]) r1--; else r2 = piv2; } else { if (nums1[piv1+1] > nums2[piv2]) r1 = piv1; else r2--; } } } } else { // remove from left // reduce array with min value on left if (!nonZero1 || !nonZero2) { // One array is empty if (nonZero1) l1 = piv1; else l2 = piv2; } else { if (piv1 > l1 && piv2 > l2) { // Both arrays have elements left of pivots if (nums1[piv1-1] < nums2[piv2-1]) l1 = piv1; else l2 = piv2; } else { if (piv1 == l1) { if (nums1[piv1] < nums2[piv2-1]) l1++; else l2 = piv2; } else { if (nums1[piv1-1] < nums2[piv2]) l1 = piv1; else l2++; } } } } } } }