Median of Two Sorted Arrays
Median of Two Sorted Arrays
Pattern: Binary Search
Idea:
Variations :
💻 Code
def find_median(A, B):
if len(A) > len(B):
A, B = B, A
m = len(A)
n = len(B)
low = 0
high = m
left_size = (m + n + 1) // 2
while low <= high:
i = (low + high) // 2
j = left_size - i
Aleft = float("-inf") if i == 0 else A[i - 1]
Aright = float("inf") if i == m else A[i]
Bleft = float("-inf") if j == 0 else B[j - 1]
Bright = float("inf") if j == n else B[j]
# Correct partition
if Aleft <= Bright and Bleft <= Aright:
if (m + n) % 2:
return max(Aleft, Bleft)
return (
max(Aleft, Bleft)
+
min(Aright, Bright)
) / 2
# Too many elements taken from A
elif Aleft > Bright:
high = i - 1
# Too few elements taken from A
else:
low = i + 1
Time complexity - O(log(min(m,n)))
Aux. Space complexity - O(1)
2 doubts : Why m+n+1 //2 and why search smaller array always? here
One of the most important Binary Search problems for FAANG interviews.
The challenge is to find the median of two sorted arrays without actually merging them.
Problem Statement
Given two sorted arrays:
A = [1, 3]
B = [2]The combined sorted order would be:
[1, 2, 3]Median:
2The obvious solution is to merge them, but the important interview solution runs in:
Approach 1: Merge
The simplest approach is to merge the two sorted arrays.
def find_median(A, B):
merged = []
i = j = 0
while i < len(A) and j < len(B):
if A[i] <= B[j]:
merged.append(A[i])
i += 1
else:
merged.append(B[j])
j += 1
merged.extend(A[i:])
merged.extend(B[j:])
n = len(merged)
if n % 2:
return merged[n // 2]
return (merged[n // 2 - 1] + merged[n // 2]) / 2Complexity
-
Time Complexity:
-
Auxiliary Space Complexity:
We can reduce the space to by merging only until we reach the median, but the time is still:
This is usually not the expected optimal solution.
Optimal Approach: Binary Search on the Partition
The key idea is:
Don’t merge the arrays. Find where to split each array so that the left half contains exactly half of all elements.
Suppose:
A = [1, 3, 8]
B = [2, 7, 10, 12]Total elements:
7We want:
Left half = 4 elements
Right half = 3 elementsWe choose a partition:
A: [1, 3 | 8]
B: [2, 7 | 10, 12]Left side contains:
1, 3, 2, 7Right side contains:
8, 10, 12If every element on the left is smaller than every element on the right, we have found the correct partition.
What Makes a Partition Correct?
Let:
A: ... | ...
B: ... | ...Define:
Aleft = largest element on A's left
Aright = smallest element on A's right
Bleft = largest element on B's left
Bright = smallest element on B's rightThe partition is correct when:
and
In other words:
largest(left side) <= smallest(right side)Once this is true, the two sides are correctly separated.
Why Do We Need Only One Binary Search?
We don’t independently choose partitions in both arrays.
If the total number of elements that must be on the left is known, then:
$$
partition_B
leftSize-partition_A
So once we choose the partition in `A`, the partition in `B` is automatically determined. Therefore we binary-search only one array. --- # Why Binary Search the Smaller Array? Always binary-search the smaller array. Suppose: ```text A has m elements B has n elements m <= n ``` Search `A`. Why? Because the partition must remain valid:0 \le partition_A \le m
O(\log m)
\boxed{O(\log(\min(m,n)))}
It also makes the partition logic safer and easier to reason about. --- # Partition Setup Let: ```text m = len(A) n = len(B) ``` Ensure: ```python if len(A) > len(B): A, B = B, A ``` Now `A` is guaranteed to be the smaller array. The number of elements that should be on the left is:\frac{m+n+1}{2}
using integer division. So: ```python left_size = (m + n + 1) // 2 ``` If we choose: ```text partition_A = i ``` then: ```text partition_B = left_size - i ``` --- # Visualizing the Partition Suppose: ```text A = [1, 3, 8] B = [2, 7, 10, 12] ``` One possible partition: ```text A = [1, 3 | 8] B = [2, 7 | 10, 12] ``` Define: ```text Aleft = 3 Aright = 8 Bleft = 7 Bright = 10 ``` Check:3 \le 10
7 \le 8
Both are true. Therefore this is the correct partition. --- # Computing the Four Boundary Values For a partition `i` in `A`: ```text Aleft = A[i - 1] Aright = A[i] ``` But what happens if the partition is at the beginning or end? Use sentinels: ```python Aleft = -inf if i == 0 Aright = inf if i == m ``` Similarly for `B`. This eliminates special-case logic. --- # Deciding Which Direction to Search Now the most important part. ## Case 1: `Aleft > Bright` ```text A: ... Aleft | Aright ... B: ... Bleft | Bright ... ``` If:Aleft > Bright
we have taken **too many elements from A** into the left half. Therefore move the partition in `A` to the **left**: ```python high = i - 1 ``` --- ## Case 2: `Bleft > Aright` We have taken **too few elements from A** into the left half. Therefore move the partition in `A` to the **right**: ```python low = i + 1 ``` --- ## Case 3: Correct Partition If:Aleft \le Bright
Bleft \le Aright
we have found the correct partition. --- # Finding the Median Once the partition is correct, we know: ```text Left side: Aleft Bleft ``` and ```text Right side: Aright Bright ``` ## Odd Number of Elements The median is the largest element on the left:\boxed{
\max(Aleft,Bleft)
}
\max(Aleft,Bleft)
\min(Aright,Bright)
\boxed{
\frac{
\max(Aleft,Bleft)+\min(Aright,Bright)
}{2}
}
\frac{3+4}{2}=3.5
--- # The Most Important Intuition Think of the two arrays as two sorted streams: ```text A: 1 3 8 9 B: 2 4 7 10 12 ``` We don't care about merging them. We only need to find the point where: ```text Everything on the LEFT ↓ is smaller than ↓ Everything on the RIGHT ``` The partition is therefore the **boundary between the lower half and upper half** of the combined sorted order. Binary Search simply moves this boundary until it becomes valid. --- # Complexity Let: ```text m = length of smaller array n = length of larger array ``` ### Time Binary search is performed only on the smaller array:\boxed{
O(\log(\min(m,n)))
}
\boxed{
O(1)
}
O(\log(\min(m,n)))
--- ### 3. Using `partition_B = (m+n)//2 - i` blindly For an odd total number of elements, using: ```text (m+n)//2 ``` can create an awkward left/right size convention. Using: ```python left_size = (m + n + 1) // 2 ``` makes the odd case particularly clean: the left side contains one extra element. --- ### 4. Forgetting empty partitions The partition can be: ```text A = [ | 1, 2, 3] ``` or: ```text A = [1, 2, 3 | ] ``` Hence the use of: ```python -inf +inf ``` for boundary values. --- ### 5. Confusing the partition index with an array index If: ```text i = 2 ``` then: ```text Aleft = A[i-1] Aright = A[i] ``` because `i` represents the **number of elements placed on the left**, not necessarily an element position. --- # Important Variations ### 1. Median of Two Sorted Arrays The classic problem. Expected:O(\log(\min(m,n)))
--- ### 2. Kth Element of Two Sorted Arrays Instead of finding the middle element, find the `k`th smallest element. The same partition idea applies. This is a very useful variation because it tests whether you actually understand the partition technique. --- ### 3. Median of Two Sorted Arrays of Different Sizes Already handled naturally by the optimal solution. The arrays do **not** need to have equal lengths. --- ### 4. Find the Kth Smallest Across Multiple Sorted Arrays The two-array partition idea can be extended, but the implementation becomes more involved. For interviews, the **two-array kth element** problem is the important one to master first. --- # Pythonic Alternative If the interviewer does **not** require the optimal algorithm, you can use: ```python import statistics statistics.median(sorted(A + B)) ``` But this is:O((m+n)\log(m+n))
O(m+n)
So these are not acceptable when the interviewer explicitly asks for the optimal solution. --- # Master Mental Model Don't memorize the entire implementation. Remember these four things: ### 1. Search the smaller array ```text A = smaller array ``` ### 2. Partition A ```text i ``` ### 3. Partition B automatically ```text j = left_size - i ``` ### 4. Check the cross-boundaries ```text Aleft <= Bright Bleft <= Aright ``` If: ```text Aleft > Bright ``` move `i` left. If: ```text Bleft > Aright ``` move `i` right. --- # Key Takeaways The entire optimal solution boils down to finding the correct partition: ```text A: [ ... left ... | ... right ... ] B: [ ... left ... | ... right ... ] ``` such that:\boxed{
Aleft \le Bright
}
\boxed{
Bleft \le Aright
}
\boxed{
median=\max(Aleft,Bleft)
}
\boxed{
median=
\frac{
\max(Aleft,Bleft)+\min(Aright,Bright)
}{2}
}
\boxed{
O(\log(\min(m,n)))\text{ time}
}
\boxed{
O(1)\text{ auxiliary space}
}