☕ DSA Java Articles

Topics › Binary Search › BS on Answers

Kth element of 2 sorted arrays

Medium TUF ID 76 SDE

Java source ↗ Edit this article ↗

Problem

You are given two arrays a (length m) and b (length n), each sorted in non-decreasing order, and an integer k with 1 ≤ k ≤ m + n. Return the element that would be at position k (1-based) if the two arrays were merged into one sorted array. Either array may be empty, and values may repeat.

Constraints: 0 ≤ m, n ≤ 10^6, 1 ≤ k ≤ m + n, values fit in a 32-bit signed integer.

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], 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 starts 72, 86, 100, 112, 113, 119, 256, ...

Approach 1: Brute force, merge into a new array

Intuition. The question is defined in terms of the merged array, so build it. The merge step from merge sort does that in linear time: keep one pointer per array and repeatedly copy the smaller front element. Then read position k - 1.

Algorithm.

  1. Allocate merged of size m + n.
  2. While both arrays have elements left, copy the smaller front element and advance that pointer.
  3. Copy whatever remains of either array.
  4. Return merged[k - 1].
static int bruteForce(int[] a, int[] b, int k) {
    int[] merged = new int[a.length + b.length];
    int i = 0, j = 0, t = 0;
    while (i < a.length && j < b.length) merged[t++] = a[i] <= b[j] ? a[i++] : b[j++];
    while (i < a.length) merged[t++] = a[i++];
    while (j < b.length) merged[t++] = b[j++];
    return merged[k - 1];
}

Complexity. Time O(m + n), space O(m + n) for the merged copy.

Approach 2: Better, merge without storing

Intuition. We never look at anything in merged except one slot, and the merge produces elements in sorted order one at a time. So we can run the same merge, count the elements as they come out, and stop at the k-th. No array is needed, and we stop as soon as we have the answer instead of finishing the whole merge.

Algorithm.

  1. Keep pointers i into a and j into b.
  2. Repeat: take from a if b is exhausted or (a is not exhausted and a[i] ≤ b[j]); otherwise take from b.
  3. Count the taken elements; when the count reaches k, return the element just taken.
static int better(int[] a, int[] b, int k) {
    int i = 0, j = 0;
    for (int count = 1; ; count++) {
        int next;
        if (j == b.length || (i < a.length && a[i] <= b[j])) next = a[i++];
        else next = b[j++];
        if (count == k) return next;
    }
}

Complexity. Time O(k), which is O(m + n) in the worst case. Space O(1).

Approach 3: Optimal, binary search on the partition

Intuition. Think about which elements make up the first k of the merged order. Because both arrays are sorted, they must be some prefix of a plus some prefix of b: if the k smallest contained a[5] but not a[4], then a[4] ≤ a[5] would also have to be among them. So the answer is determined by one number, cutA, the count taken from a; the rest, cutB = k - cutA, comes from b.

When is a split correct? Call the last taken elements leftA = a[cutA-1], leftB = b[cutB-1] and the first untaken ones rightA = a[cutA], rightB = b[cutB]. Within each array the order is already right, so the split is valid exactly when nothing taken is larger than anything left behind across the arrays:

Then the k-th element is the larger of the two last taken elements, max(leftA, leftB).

If the split is invalid, the failing condition tells us which way to move. If leftA > rightB, we took an element of a that is bigger than something we skipped in b: we took too many from a, so cutA must shrink. Otherwise leftB > rightA, and cutA must grow. Because these conditions change monotonically with cutA, we can binary search it.

The valid range of cutA is max(0, k - n) (if b cannot supply the rest) to min(k, m). Binary searching over the shorter array makes the cost O(log min(m, n)). Missing neighbours at the array ends are replaced with sentinels: a missing left element is Integer.MIN_VALUE (it never blocks), a missing right element is Integer.MAX_VALUE.

Algorithm.

  1. If a is longer than b, swap them.
  2. lo = max(0, k - n), hi = min(k, m).
  3. While lo ≤ hi: cutA = (lo + hi) / 2, cutB = k - cutA, read the four boundary values with sentinels.
    • If leftA ≤ rightB and leftB ≤ rightA, return max(leftA, leftB).
    • If leftA > rightB, set hi = cutA - 1; otherwise lo = cutA + 1.
static int optimal(int[] a, int[] b, int k) {
    if (a.length > b.length) return optimal(b, a, k);  // search over the shorter array
    int m = a.length, n = b.length;
    int lo = Math.max(0, k - n), hi = Math.min(k, m);   // cutA = how many elements we take from a
    while (lo <= hi) {
        int cutA = (lo + hi) >>> 1, cutB = k - cutA;
        int leftA = cutA == 0 ? Integer.MIN_VALUE : a[cutA - 1];
        int leftB = cutB == 0 ? Integer.MIN_VALUE : b[cutB - 1];
        int rightA = cutA == m ? Integer.MAX_VALUE : a[cutA];
        int rightB = cutB == n ? Integer.MAX_VALUE : b[cutB];
        if (leftA <= rightB && leftB <= rightA) return Math.max(leftA, leftB);
        if (leftA > rightB) hi = cutA - 1;              // took too many from a
        else lo = cutA + 1;                             // took too few from a
    }
    throw new IllegalArgumentException("k out of range");
}

Dry run on Example 2 (a is already the shorter array, m = 5, n = 7, k = 7, so lo = 0, hi = 5):

lo hi cutA cutB leftA rightA leftB rightB verdict
0 5 2 5 112 256 265 445 leftB > rightA: too few from a, lo = 3
3 5 4 3 349 770 113 119 leftA > rightB: too many from a, hi = 3
3 3 3 4 256 349 119 265 valid, answer max(256, 119) = 256

Complexity. Time O(log min(m, n)): each step halves the range of cutA in the shorter array. Space O(1).

The tests check all three methods for every k on several pairs of arrays, including duplicates, negatives and arrays of very different lengths.

Edge cases and pitfalls