门| GATE CS 2018 |第 54 题

本篇题目是 门| GATE CS 2018 |第 54 题 的介绍和解答。

题目描述

给定两个整数数组 $A, B$ 和一个数字 $k$，大小分别为 $n_A$ 和 $n_B$。找到第 $k$ 小的元素，其中有元素是从任意一个数组中取得的。

解题思路

一种较为简单的思路是将两个数组排序后合并，再取第 $k$ 小的数。这种算法的时间复杂度为 $\Theta((n_A + n_B) \log (n_A + n_B))$，空间复杂度为 $\Theta(n_A + n_B)$。

另外，由于本题中只关心第 $k$ 小的数，还有一种更快的算法：假设 $k = \frac{k}{2} + \frac{k}{2}$，我们可以比较两个数组中第 $\frac{k}{2}$ 小的数 $A[\frac{k}{2}]$ 和 $B[\frac{k}{2}]$ 的大小：

1. 如果 $A[\frac{k}{2}] < B[\frac{k}{2}]$，则说明 $A$ 中前 $\frac{k}{2}$ 个数都不可能成为第 $k$ 小的数，因此我们可以将数组 $A$ 中前 $\frac{k}{2}$ 个数舍弃掉，转而在 $A[\frac{k}{2}+1:n_A]$ 和 $B[1:n_B]$ 中寻找第 $k - \frac{k}{2}$ 小的元素。
2. 如果 $A[\frac{k}{2}] \ge B[\frac{k}{2}]$，则说明 $B$ 中前 $\frac{k}{2}$ 个数都不可能成为第 $k$ 小的数，因此可以将数组 $B$ 中前 $\frac{k}{2}$ 个数舍弃掉，转而在 $A[1:n_A]$ 和 $B[\frac{k}{2}+1:n_B]$ 中寻找第 $k - \frac{k}{2}$ 小的元素。

显然，上述过程可以通过递归实现。退出递归的条件为：

1. 数组 $A$ 或 $B$ 为空，此时直接返回另一个数组的第 $k$ 小元素；
2. $k = 1$，此时直接返回 $A[1]$ 和 $B[1]$ 中的较小值。

由于每次递归都能排除掉一个数组中的 $\frac{k}{2}$ 个数，因此每一层递归都至少要在一个数组中遍历 $\frac{k}{2}$ 个数。因此，时间复杂度为 $\Theta(\log(n_A+n_B))$。

实现代码

def kth(A, B, k):
    # 处理边界条件
    if len(A) == 0:
        return B[k-1]
    if len(B) == 0:
        return A[k-1]
    if k == 1:
        return min(A[0], B[0])
    
    # 取出两个数组中第 k/2 个元素（如果存在）
    index_A = min(k // 2, len(A))
    index_B = min(k // 2, len(B))
    item_A = A[index_A-1]
    item_B = B[index_B-1]
    
    # 如果 A[k/2] < B[k/2]，则令 k = k - k/2，否则令 k = k - k/2
    if item_A < item_B:
        return kth(A[index_A:], B, k - index_A)
    else:
        return kth(A, B[index_B:], k - index_B)