寻找两个有序数组交叉和中第k个元素的最快算法

假设两个浮点数组x[]和y[]均为升序排列（若为降序可先反转处理），交叉和数组s包含所有x[i] + y[j]的组合，我们需要找到s中的第k小元素，且k远小于min(n, m)（n、m分别为两个数组的长度）。

最优方法：最小堆（优先队列） #

针对k极小的场景，这是效率最高的方案之一，时间复杂度为O(k log k)，空间复杂度为O(k)，远优于你当前的实现。

核心思路 #

利用最小堆维护当前待选的最小和，每次取出堆顶的最小元素后，仅扩展其相邻的两个新组合（x[i+1]+y[j]和x[i]+y[j+1]），同时用集合记录已访问过的(i,j)索引对，避免重复加入堆。

步骤 #

1. 初始化最小堆，将初始组合(x[0]+y[0], 0, 0)加入堆，并用集合标记(0,0)为已访问。
2. 重复k次以下操作：
   • 弹出堆顶元素，该元素即为当前第t小的和（t从1到k）。
   • 若i+1 < n且(i+1,j)未被访问，计算x[i+1]+y[j]并加入堆，标记该索引对为已访问。
   • 若j+1 < m且(i,j+1)未被访问，计算x[i]+y[j+1]并加入堆，标记该索引对为已访问。
3. 第k次弹出的元素即为目标的第k小和。

代码示例（Python） #

import heapq

def find_kth_smallest_sum(x, y, k):
    n, m = len(x), len(y)
    visited = set()
    heap = []
    
    initial_sum = x[0] + y[0]
    heapq.heappush(heap, (initial_sum, 0, 0))
    visited.add((0, 0))
    
    result = 0
    for _ in range(k):
        result, i, j = heapq.heappop(heap)
        # 扩展x的下一个元素与当前y的组合
        if i + 1 < n and (i+1, j) not in visited:
            heapq.heappush(heap, (x[i+1] + y[j], i+1, j))
            visited.add((i+1, j))
        # 扩展当前x与y的下一个元素的组合
        if j + 1 < m and (i, j+1) not in visited:
            heapq.heappush(heap, (x[i] + y[j+1], i, j+1))
            visited.add((i, j+1))
    return result

替代方案：二分查找 #

若k不是极小但仍远小于min(n,m)，可采用二分查找法，时间复杂度为O(k log k + log(max_diff))（max_diff为二分范围的差值），需注意浮点精度问题：

1. 确定二分范围：左边界为x[0]+y[0]，右边界为x[k-1]+y[k-1]（第k小元素必然不超过该值）。
2. 对中间值mid，用双指针法统计有多少个x[i]+y[j] <= mid（仅需遍历x的前k个元素，对每个x[i]用二分找y中符合条件的最大索引）。
3. 根据计数调整边界：若计数≥k，说明第k小元素≤mid，缩小右边界；否则扩大左边界。
4. 当边界收敛到足够小的精度范围内，左边界即为目标值。

你当前方法的问题 #

1. 重复元素冗余：x[i]+y[j]与y[j]+x[i]是同一值，重复生成会浪费计算和存储资源。
2. 不必要的元素生成：第k小元素必然来自x的前k个元素与y的前k个元素的组合（因数组升序，超出该范围的组合和必然更大），无需生成x前k个与所有y、y前k个与所有x的组合。即使仅生成前k×k个组合，其数量仍远大于k，效率远低于堆方法。

内容的提问来源于stack exchange，提问作者user2961927