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)).

You may assume nums1 and nums2 cannot be both empty.

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

Problem Summary #

Given two sorted arrays nums1 and nums2 of sizes m and n.

Find the median of these two sorted arrays, and the time complexity of the algorithm must be O(log(m + n)).

You may assume nums1 and nums2 will not both be empty.

Solution Ideas #

Given two sorted arrays, find the median in the sorted array after merging these two arrays. The required time complexity is O(log (m+n)).
The easiest approach that comes to mind for this problem is to merge the two arrays and then take the median. However, merging sorted arrays is an O(m+n) operation, which does not meet the requirement. Seeing the log time complexity given in the problem naturally suggests binary search.
Since we need to find the median of the final merged array, and the total size of the two arrays is also known, the middle position is known as well. We only need to binary search the partition position in one array, and the partition position in the other array can also be obtained. To minimize the time complexity, binary search the shorter of the two arrays.
The key question is how to partition array 1 and array 2. In fact, it is how to partition array 1. First randomly generate a midA by binary search. When does the partition line satisfy the median condition? That is, all numbers on the left side of the line are less than the numbers on the right side, namely, nums1[midA-1] ≤ nums2[midB] && nums2[midB-1] ≤ nums1[midA]. If these conditions are not satisfied, the partition line needs to be adjusted. If nums1[midA] < nums2[midB-1], it means the numbers on the left side divided by this midA line are too small, and the partition line should move right; if nums1[midA-1] > nums2[midB], it means the numbers on the left side divided by this midA line are too large, and the partition line should move left. After multiple adjustments, a partition line satisfying the conditions can always be found.
Suppose the two partition lines have now been found. The indices on both sides of the partition line in array 1 are midA - 1 and midA, respectively. The indices on both sides of the partition line in array 2 are midB - 1 and midB, respectively. After finally merging into the final array, if the array length is odd, then the median is max(nums1[midA-1], nums2[midB-1]). If the array length is even, then the two numbers in the middle positions are, in order: max(nums1[midA-1], nums2[midB-1]) and min(nums1[midA], nums2[midB]), so the median is (max(nums1[midA-1], nums2[midB-1]) + min(nums1[midA], nums2[midB])) / 2. See the figure below:

Code #


package leetcode

func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
	// Assume nums1 is shorter
	if len(nums1) > len(nums2) {
		return findMedianSortedArrays(nums2, nums1)
	}
	low, high, k, nums1Mid, nums2Mid := 0, len(nums1), (len(nums1)+len(nums2)+1)>>1, 0, 0
	for low <= high {
		// nums1:  ……………… nums1[nums1Mid-1] | nums1[nums1Mid] ……………………
		// nums2:  ……………… nums2[nums2Mid-1] | nums2[nums2Mid] ……………………
		nums1Mid = low + (high-low)>>1 // The right side of the partition is mid, and the left side is mid - 1
		nums2Mid = k - nums1Mid
		if nums1Mid > 0 && nums1[nums1Mid-1] > nums2[nums2Mid] { // The partition in nums1 is too far right; move left
			high = nums1Mid - 1
		} else if nums1Mid != len(nums1) && nums1[nums1Mid] < nums2[nums2Mid-1] { // The partition in nums1 is too far left; move right
			low = nums1Mid + 1
		} else {
			// Found a suitable partition; need to output the final result
			// There are 2 cases: odd and even
			break
		}
	}
	midLeft, midRight := 0, 0
	if nums1Mid == 0 {
		midLeft = nums2[nums2Mid-1]
	} else if nums2Mid == 0 {
		midLeft = nums1[nums1Mid-1]
	} else {
		midLeft = max(nums1[nums1Mid-1], nums2[nums2Mid-1])
	}
	if (len(nums1)+len(nums2))&1 == 1 {
		return float64(midLeft)
	}
	if nums1Mid == len(nums1) {
		midRight = nums2[nums2Mid]
	} else if nums2Mid == len(nums2) {
		midRight = nums1[nums1Mid]
	} else {
		midRight = min(nums1[nums1Mid], nums2[nums2Mid])
	}
	return float64(midLeft+midRight) / 2
}