Kth element of 2 sorted arrays

Medium TUF ID 2767 A2Z

Problem

You are given two sorted arrays a (length n) and b (length m) and an integer k with 1 ≤ k ≤ n + m. If the two arrays were merged into one sorted array, which value would be at position k (1-indexed)? Return that value.

Constraints: 0 ≤ n, m ≤ 10^5, n + m ≥ 1, values fit in a 32-bit signed integer. Duplicates may appear within and across the arrays.

Example 1

a = [2, 3, 6, 7, 9], b = [1, 4, 8, 10], k = 5
answer = 6           // merged = [1, 2, 3, 4, 6, 7, 8, 9, 10], the 5th element is 6

Example 2

a = [100, 112, 256, 349, 770], b = [72, 86, 113, 119, 265, 445, 892], k = 7
answer = 256         // merged = [72, 86, 100, 112, 113, 119, 256, ...]

Example 3

a = [], b = [1, 2, 3], k = 2
answer = 2           // with one array empty, the answer is simply b[k - 1]

Approach 1: Brute force, merge and index

Intuition. The statement describes the merged array, so build it. The classic two-pointer merge compares the heads of the two arrays, copies the smaller one, and advances that pointer. Once the merged array exists, the answer is at index k − 1.

Algorithm.

Allocate merged of size n + m.
While both arrays still have elements, copy the smaller head; then copy the leftovers of whichever array remains.
Return merged[k − 1].

static int bruteForce(int[] a, int[] b, int k) {
    int n = a.length, m = b.length;
    int[] merged = new int[n + m];
    int i = 0, j = 0, p = 0;
    while (i < n && j < m) merged[p++] = a[i] <= b[j] ? a[i++] : b[j++];
    while (i < n) merged[p++] = a[i++];
    while (j < m) merged[p++] = b[j++];
    return merged[k - 1];
}

Complexity. Time O(n + m) for the full merge. Space O(n + m) for the copy.

Approach 2: Better, merge only as far as needed

Intuition. Merging past position k is wasted work, and storing anything is wasted memory. Run the same merge but keep only a counter; the k-th element produced is the answer, so return it immediately. The merge never has to look at elements beyond that point.

Algorithm.

Keep pointers i, j and a counter taken.
Repeatedly produce the next element of the merge: take from a if b is exhausted or a’s head is not larger, otherwise from b.
Increment taken; when it reaches k, return the element just produced.

static int better(int[] a, int[] b, int k) {
    int n = a.length, m = b.length;
    int i = 0, j = 0, taken = 0;
    while (true) {
        int next;
        if (j >= m || (i < n && a[i] <= b[j])) next = a[i++]; else next = b[j++];
        taken++;
        if (taken == k) return next;
    }
}

The condition is ordered so that no array is read out of bounds: j >= m is tested before b[j], and i < n before a[i]. Because k ≤ n + m, the loop always returns before both arrays run dry.

Complexity. Time O(k), at most O(n + m). Space O(1).

This is enough for many inputs, but when k is around the middle of two large arrays it is still linear. The optimal solution stops walking altogether.

Approach 3: Optimal, binary search on the cut

Intuition. The k smallest elements of the merged array consist of some prefix of a and some prefix of b: say the first i elements of a and the first j = k − i elements of b. If we knew i, the answer would be the larger of a[i-1] and b[j-1], the largest element among those k.

How do we recognise the right i? The chosen prefixes form the k smallest elements exactly when nothing in the chosen parts is larger than anything in the unchosen parts. Within each array that holds automatically (they are sorted), so only two cross checks remain: a[i-1] <= b[j] and b[j-1] <= a[i]. If a[i-1] > b[j], we took too many from a and should lower i; if b[j-1] > a[i], we took too few and should raise i. These tests move in one direction as i grows, so binary search over i works.

The range of i needs a little care: a can contribute at most min(k, n) elements, and at least k − m (when b contributes everything it has). Searching over the shorter array keeps the number of probes to O(log(min(n, m))).

Algorithm.

If a is longer than b, swap them.
Set lo = max(0, k − m) and hi = min(k, n).
While lo <= hi: let i = mid, j = k − i.
- l1 = a[i-1] or −∞ if i = 0; l2 = b[j-1] or −∞ if j = 0; r1 = a[i] or +∞ if i = n; r2 = b[j] or +∞ if j = m.
- If l1 <= r2 and l2 <= r1, return max(l1, l2).
- If l1 > r2, set hi = i − 1; otherwise set lo = i + 1.

static int optimal(int[] a, int[] b, int k) {
    if (a.length > b.length) return optimal(b, a, k);  // always search over the shorter array
    int n = a.length, m = b.length;
    int lo = Math.max(0, k - m), hi = Math.min(k, n);   // a contributes at least k-m and at most min(k, n) elements
    while (lo <= hi) {
        int i = lo + (hi - lo) / 2;                      // i from a and j = k-i from b form the k smallest
        int j = k - i;
        int l1 = i == 0 ? Integer.MIN_VALUE : a[i - 1];
        int l2 = j == 0 ? Integer.MIN_VALUE : b[j - 1];
        int r1 = i == n ? Integer.MAX_VALUE : a[i];
        int r2 = j == m ? Integer.MAX_VALUE : b[j];
        if (l1 <= r2 && l2 <= r1) return Math.max(l1, l2);
        if (l1 > r2) hi = i - 1;                         // too many from a: move the cut left
        else lo = i + 1;                                 // too few from a: move the cut right
    }
    throw new IllegalArgumentException("inputs are not sorted or k is out of range");
}

Dry run on a = [2, 3, 6, 7, 9], b = [1, 4, 8, 10], k = 5. a is longer, so the arrays swap: search over a = [1, 4, 8, 10] (n = 4) against b = [2, 3, 6, 7, 9] (m = 5). lo = max(0, 5 − 5) = 0, hi = min(5, 4) = 4.

lo	hi	i	j	l1	l2	r1	r2	l1 ≤ r2 and l2 ≤ r1?	result
0	4	2	3	4	6	8	7	yes (4 ≤ 7, 6 ≤ 8)	max(4, 6) = 6

The first probe already lands on the valid cut: the five smallest are {1, 4} from one array and {2, 3, 6} from the other, and the largest among them is 6.

Complexity. Time O(log(min(n, m))): each probe does constant work and halves the range of i. Space O(1).

Why must a valid cut exist in [lo, hi]? Take the real merged array and look at its first k elements; they consist of some prefix of a of length i* and a prefix of b of length k − i*. That i* lies within the computed bounds and satisfies both cross checks. With duplicates several i may be valid, but all of them give the same max(l1, l2), because the k-th smallest value is a property of the multiset, not of the cut.

Edge cases and pitfalls

One array empty. With n = 0 the range collapses to i = 0, j = k, and the sentinel l1 = −∞, r1 = +∞ lets the answer b[k-1] fall out of the general formula.
k = 1 and k = n + m. These return the global minimum and maximum. The bounds lo = max(0, k − m) and hi = min(k, n) keep j legal in both extremes; forgetting k − m as a lower bound would let j exceed m.
Sentinels are also legal values. Integer.MIN_VALUE and Integer.MAX_VALUE stand in for −∞ and +∞, but since the checks use <=, real extreme values in the arrays still compare correctly.
Searching over the longer array. The swap at the top is not cosmetic: it is what guarantees j = k − i stays inside [0, m] for every i in the range, and what makes the complexity depend on the shorter array.
Duplicates. [1, 1, 1] with [1, 1] and any k returns 1; equal elements pass both <= checks.
Off-by-one in k. The problem is 1-indexed. The brute force reads merged[k − 1], and the optimal method puts exactly k elements on the left.

Median of 2 sorted arrays is the same cut search with k fixed at the middle, plus the even-length average.
Merge two sorted arrays without extra space is the merge from the brute force done in place.
Kth smallest element in an array answers the same question for a single unsorted array, using a heap.
Merge K sorted Lists generalises the merge to many inputs.

Kth element of 2 sorted arrays

Problem

Approach 1: Brute force, merge and index

Approach 2: Better, merge only as far as needed

Approach 3: Optimal, binary search on the cut

Edge cases and pitfalls

Related problems