从多个有序数组中提取Top K最大元素的最高效方法

首先明确前提：题中给出的数组均为降序排列（从大到小），这种场景下效率最高的方案是「最大堆+多指针遍历」，整体时间复杂度为O(M + K*logM)，其中M是输入数组的总个数，K是需要取的最大值数量。这个方案不需要遍历所有数组元素，也不需要对单个数组做额外排序，取够K个元素就可以终止，额外空间复杂度只有O(M)，是当前场景下的理论最优解。

核心思路 #

每个降序排列的数组里，未被取用的元素的最大值一定在当前遍历位置的最前端，我们只需要每次从所有数组的当前可用最大值里挑出全局最大的那个，放进结果集，再把该元素所在数组的下一个元素加入候选池即可，重复这个流程直到拿满K个元素或者所有元素被取完。
用最大堆来维护候选池，可以把每次找全局最大值的开销从O(M)降到O(logM)，效率提升非常明显。

流程演示（对应题目示例） #

题目输入：

• K = 5
• array1: [100, 50, 30, 20, 10]
• array2: [400, 50, 20, 0]
• array3: [0, 0, 0, 0]

执行流程：

1. 初始化堆，把每个数组的第一个元素（各自的当前最大值）入堆，堆内元素为400、100、0
2. 弹出堆顶最大值400加入结果，再把array2的下一个元素50入堆，当前结果[400]，堆内元素100、50、0
3. 弹出堆顶最大值100加入结果，再把array1的下一个元素50入堆，当前结果[400,100]，堆内元素50、50、0
4. 弹出堆顶最大值50（来自array2）加入结果，再把array2的下一个元素20入堆，当前结果[400,100,50]，堆内元素50、20、0
5. 弹出堆顶最大值50（来自array1）加入结果，再把array1的下一个元素30入堆，当前结果[400,100,50,50]，堆内元素30、20、0
6. 弹出堆顶最大值30加入结果，此时结果长度达到5，直接返回[400, 100, 50, 50, 30]，和示例输出完全一致。

参考实现（Python） #

import heapq

def get_top_k(arrays, k):
    # Python内置heapq是最小堆，存入元素负值模拟最大堆
    heap = []
    result = []
    # 初始化堆：存入每个非空数组的首元素、所属数组索引、元素位置
    for arr_idx, arr in enumerate(arrays):
        if arr:
            heapq.heappush(heap, (-arr[0], arr_idx, 0))
    
    while len(result) < k and heap:
        neg_val, arr_idx, cur_pos = heapq.heappop(heap)
        result.append(-neg_val)
        # 如果当前数组还有后续元素，将下一个元素加入候选堆
        next_pos = cur_pos + 1
        if next_pos < len(arrays[arr_idx]):
            next_val = arrays[arr_idx][next_pos]
            heapq.heappush(heap, (-next_val, arr_idx, next_pos))
    return result

其他方案的劣势 #

• 全量合并排序：需要遍历所有元素，排序开销为O(T log T)（T为所有数组总元素数），当T远大于K时会产生大量无意义的计算开销
• 线性扫描选最大值：每次遍历所有数组的当前指针位置找最大值，单次取数开销为O(M)，总时间复杂度O(K*M)，当数组数量M较大时，效率远低于堆方案

如果输入的数组是升序排列，只需要把初始指针放到数组末尾，从后往前取元素即可，核心逻辑完全不变。

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