GeeksforGeeks Hard

Kth Element of Two Sorted Arrays

Two sorted arrays; return the k-th smallest of their union in logarithmic time.

binary-searcharray

02

Intuition

Take exactly k elements as a prefix of A plus a prefix of B. The split is valid when nothing left of either cut exceeds anything right of the other cut. Binary search how many of the k come from A — validity is monotone in that count.

How to spot this pattern

Median-of-two-sorted-arrays generalised: binary search the partition so that exactly k elements fall on the left. The bounds max(0, k - m) and min(k, n) matter — they keep j = k - i inside the second array. Searching a partition point rather than a value is the transferable move.

03

Approach

1

Partition, don't merge

Choose i from A and k−i from B. The k-th smallest is max(A[i−1], B[k−i−1]) for the valid split — no O(k) walk.

2

Validity check

Split is valid iff A[i−1] ≤ B[k−i] and B[k−i−1] ≤ A[i] (missing neighbours = ±∞). If A's last taken is too big, take fewer from A.

3

Search on the smaller array

Binary search i in [max(0,k−m), min(k,n)] on the shorter array → O(log min(n,m,k)).

04

Solution & live demo

▶1def kth_element(a, b, k):
▶2    if len(a) > len(b): a, b = b, a
▶3    n, m = len(a), len(b)
▶4    lo, hi = max(0, k - m), min(k, n)
▶5    INF = float("inf")
▶6    while lo <= hi:
▶7        i = (lo + hi) // 2          # take i from a
▶8        j = k - i                   # and j from b
▶9        a_left  = a[i-1] if i > 0 else -INF
▶10        a_right = a[i]   if i < n else INF
▶11        b_left  = b[j-1] if j > 0 else -INF
▶12        b_right = b[j]   if j < m else INF
▶13        if a_left <= b_right and b_left <= a_right:
▶14            return max(a_left, b_left)
▶15        if a_left > b_right: hi = i - 1
▶16        else: lo = i + 1

05

Common pitfalls

Merging until the k-th element

✗ Wrong

merged = sorted(a + b)
return merged[k - 1]

✓ Right

i = (lo + hi) // 2
j = k - i

O((n + m) log(n + m)) to read one element. Binary searching the split touches only the four values straddling it, giving O(log min(n, m)).

Using 0 and n as the search bounds

✗ Wrong

lo, hi = 0, n

✓ Right

lo, hi = max(0, k - m), min(k, n)

j = k - i must be a legal index into b. If k exceeds m, taking too few from a pushes j past the end of b; the tightened bounds make every probe valid without extra guards.

Binary searching the longer array

✗ Wrong

# no swap
n, m = len(a), len(b)

✓ Right

if len(a) > len(b): a, b = b, a

Complexity is O(log min(n, m)) only when you search the shorter side, and the bound arithmetic assumes it. Swapping up front costs nothing and keeps both properties.

06

Edge cases

k larger than one array

Lower bound max(0, k−m) forces enough elements from the other array.

All of one array first, e.g. A=[1,2], B=[10,20], k=2

Valid split takes i=2, j=0; sentinels ±∞ make the checks pass.

07

Complexity

Time

O(log min(n, m, k))

Space

O(1)

Pure partition search; no merging.