Find the k-th Smallest Element in the Union of Two Sorted Arrays

时间:2022-01-01 22:55:11

(http://leetcode.com/2011/01/find-k-th-smallest-element-in-union-of.html)

Given two sorted arrays A, B of size m and n respectively. Find the k-th smallest element in the union of A and B. You can assume that there are no duplicate elements.

O(lg m + lg n) solution:

1. Maintaining the invariant

i + j = k - 1,

2. If Bj-1 < A< Bj, then Ai must be the k-th smallest,

or else if Ai-1 < B< Ai, then Bj must be the k-th smallest.

code:

int findKthSmallest(int A[], int m, int B[], int n, int k)
{
    assert(A && m >= 0 && B && n >= 0 && k > 0 && k <= m+n);

    int i = (int)((double)m / (m+n) * (k-1));
    int j = (k-1) - i;

    assert(i >= 0 && j >= 0 && i <= m && j <= n);

    int Ai_1 = ((i == 0) ? INT_MIN : A[i-1]);
    int Bj_1 = ((j == 0) ? INT_MIN : b[j-1]);
    int Ai = ((i == m) ? INT_MAX : A[i]);
    int Bj = ((j == n) ? INT_MAX : B[j]);

    if (Bj_1 < Ai && Ai < Bj)
        return Ai;
    else if (Ai_1 < Bj && Bj < Ai)
        return Bj;

    assert((Ai > Bj && Ai_1 > Bi) || (Ai < Bj && Ai < Bj_1));

    if (Ai < Bj)
        return findKthSmallest(A+i+1, m-i-1, B, j, k-i-1);
    else
        return findKthSmallest(A, i, B+j+1, n-j-1, k-j-1);
}