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.

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

python

▶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

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.

06

Complexity

Time

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

Space

O(1)

Pure partition search; no merging.