多有序数组的第K小元素求解：如何扩展双数组二分法？

问题回顾 #

你已经掌握了两个有序数组中寻找第k小元素的O(log(min(m,n)))二分算法，但无法扩展到三个及以上有序数组的场景（比如示例中的array1 = [2,3,6,7,9]、array2 = [1,4,8,10]、array3 = [2,3,5,7]），下面给出可行的解决方案，并说明时间复杂度情况。

可行解决方案 #

方法1：全局二分候选值（最通用） #

这是处理多有序数组第k小问题的标准方法，核心是通过二分法猜测候选值，统计所有数组中小于等于该值的元素总数，以此缩小范围找到目标值：

• 步骤：
  1. 确定二分的上下界：取所有数组的最小值作为low，最大值作为high。
  2. 迭代二分：
     • 计算中间值mid = (low + high) // 2
     • 遍历每个数组，用二分法统计其中小于等于mid的元素个数，求和得到count
     • 若count < k：说明第k小元素比mid大，将low调整为mid + 1
     • 若count >= k：说明第k小元素小于等于mid，将high调整为mid
  3. 终止条件：当low == high时，该值就是第k小元素。

代码实现

import bisect

def find_kth_smallest(arrays, k):
    # 初始化二分边界（跳过空数组）
    valid_arrays = [arr for arr in arrays if arr]
    low = min(arr[0] for arr in valid_arrays)
    high = max(arr[-1] for arr in valid_arrays)
    
    while low < high:
        mid = (low + high) // 2
        total = 0
        for arr in valid_arrays:
            # 统计当前数组中<=mid的元素数量
            cnt = bisect.bisect_right(arr, mid)
            total += cnt
            # 提前终止，避免多余计算
            if total >= k:
                break
        if total < k:
            low = mid + 1
        else:
            high = mid
    return low

# 测试示例
array1 = [2,3,6,7,9]
array2 = [1,4,8,10]
array3 = [2,3,5,7]
print(find_kth_smallest([array1, array2, array3], 5))  # 输出：4

方法2：堆解法（适合k远小于总元素数的场景） #

如果k的值远小于所有数组的总元素数，用堆的方法效率更高：

• 步骤：
  1. 维护一个最小堆，初始时将每个数组的第一个元素及其所在数组的索引、元素位置存入堆。
  2. 循环弹出堆顶元素k次：
     • 每次弹出后，若该元素所在数组还有下一个元素，就将下一个元素加入堆。
  3. 第k次弹出的元素就是目标值。

代码实现

import heapq

def find_kth_smallest_heap(arrays, k):
    heap = []
    # 初始化堆：(元素值, 数组索引, 元素在数组中的位置)
    for i, arr in enumerate(arrays):
        if arr:
            heapq.heappush(heap, (arr[0], i, 0))
    
    count = 0
    result = 0
    while heap:
        val, arr_idx, elem_idx = heapq.heappop(heap)
        count += 1
        if count == k:
            result = val
            break
        # 加入当前数组的下一个元素
        if elem_idx + 1 < len(arrays[arr_idx]):
            next_val = arrays[arr_idx][elem_idx + 1]
            heapq.heappush(heap, (next_val, arr_idx, elem_idx + 1))
    return result

# 测试示例
print(find_kth_smallest_heap([array1, array2, array3], 5))  # 输出：4

时间复杂度说明 #

无法达到双数组场景的O(log(min(m,n)))复杂度，原因如下：

• 双数组的算法依赖两个数组切割点的严格约束（cut1 + cut2 = k），能通过单次二分直接定位目标；但多数组场景下，无法建立如此简洁的切割点约束关系。
• 全局二分法的时间复杂度为O(t * log(L) * log(R))：其中t是数组数量，L是单个数组的最大长度，R是所有数组元素的最大值与最小值的差。
• 堆解法的时间复杂度为O(k log t)：当k远小于总元素数时，这个复杂度会优于全局二分法。

内容的提问来源于stack exchange，提问作者M.Soyturk