两个有序数组中的第 k 小元素

介绍

存在两个非降序排列的整数数组 $A$ 和 $B$，数组的长度分别为 $m$ 和 $n$，找出它们合并后的数组中的第 $k$ （从 $1$ 开始）个数。

1
2


def findKth(A: List[int], B: List[int], k: int) -> int:
    ...

1
2


int find_kth(vector<int> &A, vector<int> &B, int k) {
}

解决方法

遍历

最基本的方法就是按照 合并两个有序数组 的步骤来寻找第 $k$ 个数。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def find_kth(A: List[int], B: List[int], k: int) -> int:
    m, n = len(A), len(B)
    i = j = ans = 0

    while i + j < k:
        if j == n or (i < m and A[i] <= B[j]):
            ans = A[i]
            i += 1
        else:
            ans = B[j]
            j += 1

    return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


int find_kth(vector<int> &A, vector<int> &B, int k) {
    int m = A.size(), n = B.size();
    int i = 0, j = 0, ans = 0;

    while (i + j < k) {
        if (j == n || (i < m && A[i] <= B[j])) {
            ans = A[i++];
        } else {
            ans = B[j++];
        }
    }

    return ans;
}

时间复杂度：$\mathcal{O}(k)$
空间复杂度：$\mathcal{O}(1)$

二分搜索

方法一

我们首先来判断第 $k$ 个数能否在数组 $B$ 中。

如果 $B[i]$ 可以排在第 $k$ 个数，那么必须最多（$\le$）有 $k-1$ 个数小于 $B[i]$，至少（$\ge$）有 $k-1$ 个数 小于等于 $B[i]$。即如果小于 $B[i]$ 的数的 个数多于（$\gt$）$k-1$ 个，那么 $B[i]$ 必然排在第 $k$ 个数之后。如果 小于等于 $B[i]$ 的数的 个数少于（$\lt$）$k-1$ 个，那么 $B[i]$ 必然排在第 $k$ 个数之前。

假设 $B[i]$ 按顺序可以插在 $A$ 中的 $[a,b]$ 位置，即 C++ 中 equal_range 表示的范围，如 $3$ 可以插在数组 $[1,2,2,3,3,3,4,5]$ 中的下标位置为 $3,4,5,6$。

即范围 $[a,b]$ 说明了 $A$ 中下标范围为 $[0,a-1]$ 的数都小于 $B[i]$，共 $a$ 个。下标范围为 $[0,b-1]$ 的数都 小于等于 $B[i]$，共 $b$ 个。

而在 $B$ 中下标范围为 $[0,i-1]$ 的数都 小于等于 $B[i]$，共 $i$ 个。

所以 $B[i]$ 可以排列的位置为 $[a+i+1,b+i+1]$。

如果 $k \lt a+i+1$，说明 $B[i]$ 排在第 $k$ 个数之后，第 $k$ 个数可能在 $[0,i-1]$ 中。
如果 $k \gt b+i+1$，说明 $B[i]$ 排在第 $k$ 个数之前，第 $k$ 个数可能在 $[i+1,n-1]$ 中。
否则 $a+i+1 \le k \le b+i+1$，$B[i]$ 一定可以可以排在第 $k$ 个数位置，直接返回 $B[i]$ 就可以。

如果上面的步骤找不到 $B[i]$，则说明 $B$ 中不存在可以排在第 $k$ 个位置的数，那第 $k$ 个数一定在 $A$ 中，此时直接返回 find_kth(B, A, k) 重新进行搜索就可以。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


def find_kth(A: List[int], B: List[int], k: int) -> int:
    L, R = 0, len(B) - 1

    while L <= R:
        i = L + ((R - L) >> 1)
        a, b = bisect_left(A, B[i]), bisect_right(A, B[i])

        if k < a + i + 1:
            R = i - 1
        elif k > b + i + 1:
            L = i + 1
        else:
            return B[i]

    return findKth(B, A, k)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


int find_kth(vector<int> &A, vector<int> &B, int k) {
    int L = 0, R = B.size() - 1;

    while (L <= R) {
        int i = L + ((R - L) >> 1);
        auto range = equal_range(A.begin(), A.end(), B[i]);
        int a = range.first - A.begin(), b = range.second - A.begin();

        if (k < a + i + 1) {
            R = i - 1;
        } else if (k > b + i + 1) {
            L = i + 1;
        } else {
            return B[i];
        }
    }

    return find_kth(B, A, k);
}

时间复杂度：$\mathcal{O}(\log(m)\log(n))$
空间复杂度：$\mathcal{O}(1)$