算法测验 | SP2竞赛1 | 问题4

本次算法测验针对SP2竞赛的第一道难度中的第四个问题，旨在考察参赛者对排序算法的理解和实现能力。

问题描述

给定一个长度为 $n$ 的数组 $a$，请你设计一个算法，找到数组中前 $k$ 大的数，并按从大到小的顺序输出它们。

解题思路

本题可以用堆排序的思想来解决。首先构建一个大小为 $k$ 的小根堆，用 $a_1$ 到 $a_k$ 初始化堆。然后从 $a_{k+1}$ 开始，将后面的数依次与堆顶元素比较。若后面的数比堆顶元素大，则将堆顶元素删除并将该数插入堆，否则不做任何处理。最终堆中剩下的就是前 $k$ 大的数。

具体实现可以分为两个步骤：初始化堆和依次处理后面的数。对于初始化堆，可以用 $O(k)$ 的时间复杂度完成。对于依次处理后面的数，则需要匹配 $n-k$ 个数，每次操作的时间复杂度为 $O(\log k)$，所以总的时间复杂度为 $O(k+(n-k)\log k)=O(n\log k)$。空间复杂度为 $O(k)$。

代码实现

下面是用 Python 实现的代码片段：

import heapq

def largest_k_numbers(a, k):
    heap = a[:k]
    heapq.heapify(heap)
    for i in range(k, len(a)):
        if a[i] > heap[0]:
            heapq.heappop(heap)
            heapq.heappush(heap, a[i])
    return sorted(heap, reverse=True)

需要注意的是，这里用了 Python 标准库中的堆实现 heapq。因为 Python 中没有直接封装大根堆的库，所以我们可以用小根堆反过来实现大根堆。当然，用其他语言实现同样的算法也是可行的。