Find the median of two sorted arrays in log time. Binary search a cut in the shorter one so both left halves hold the smaller half.
▼The problem
LeetCode 4 (Hard). Given two sorted arrays nums1 (length m) and nums2 (length n), return the median of the two arrays together. The overall run time should be O(log(m+n)).
Examples (LeetCode's): [1,3], [2] → 2.0 (merged 1 2 3), and [1,2], [3,4] → 2.5 (merged 1 2 3 4, the middle pair averaged).
The solution
from math import inf
def findMedianSortedArrays(a, b):
if len(a) > len(b): a, b = b, a
m, n = len(a), len(b)
half = (m + n + 1) // 2
lo, hi = 0, m
while lo <= hi:
i = (lo + hi) // 2
j = half - i
aL = a[i-1] if i > 0 else -inf
aR = a[i] if i < m else inf
bL = b[j-1] if j > 0 else -inf
bR = b[j] if j < n else inf
if aL <= bR and bL <= aR:
if (m + n) % 2: return max(aL, bL)
return (max(aL, bL) + min(aR, bR)) / 2
if aL > bR: hi = i - 1
else: lo = i + 1Transcript
Median of Two Sorted Arrays. Given two sorted arrays, return the median of all their numbers, in order log of m plus n time.
One three, and two: together that's one, two, three, so the median is two. One two, and three four: an even count, so average the middle pair: two point five.
The simple way merges both arrays and takes the middle. That's order m plus n time, and it ignores that they're already sorted.
Picture a tug of war. Both arrays stand in rows, and a center line splits everyone into two teams, with half the players on the left. Choose the cut in the shorter row, and the other row's cut follows. It's fair when each row's last left player is at most the other row's first right player. If A's is too big, move its cut left. Otherwise, move it right.
In code, binary search the cut i in the shorter array, A. J is half minus i. Read the four border values, with infinity past an end. If both checks pass, the median is the max of the lefts, averaged with the min of the rights when the count is even.
Try one, three, eight, nine, fifteen and four, six, seven, twelve, eighteen, twenty one. The left team needs six. Cut A after two: twelve beats eight, so move right. Cut after four: nine beats seven, so move left. Cut after three: both checks pass, so the median is eight.
Each try halves the range: order log of min of m and n time, and constant space.
Cut the short row, balance the teams, check the corners. That's Median of Two Sorted Arrays.