跳转至

两个有序数组中的第 k 小元素#
发布于2021-02-20
上次编辑2021-04-20

Introduction

存在两个非降序排列的整数数组 \(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\) 中。

Info

如果 \(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)\)

返回顶部

在手机上阅读