C++Top-K 問題

Question

給定一個長度為 n 無序數(shù)組 nums ，請返回數(shù)組中前 k 大的元素。

對于該問題，我們先介紹兩種思路比較直接的解法，再介紹效率更高的堆解法。

8.3.1   方法一：遍歷選擇

我們可以進(jìn)行圖 8-6 所示的 k 輪遍歷，分別在每輪中提取第 1、2、…、k 大的元素，時間復(fù)雜度為 O(nk) 。

此方法只適用于 k?n 的情況，因為當(dāng) k 與 n 比較接近時，其時間復(fù)雜度趨向于 O(n2) ，非常耗時。

圖 8-6   遍歷尋找最大的 k 個元素

Tip

當(dāng) k=n 時，我們可以得到完整的有序序列，此時等價于“選擇排序”算法。

方法二：排序

如圖 8-7 所示，我們可以先對數(shù)組 nums 進(jìn)行排序，再返回最右邊的 k 個元素，時間復(fù)雜度為 O(nlog?n) 。

顯然，該方法“超額”完成任務(wù)了，因為我們只需要找出最大的 k 個元素即可，而不需要排序其他元素。

圖 8-7   排序?qū)ふ易畲蟮?k 個元素

方法三：堆

我們可以基于堆更加高效地解決 Top-K 問題，流程如圖 8-8 所示。

1. 初始化一個小頂堆，其堆頂元素最小。
2. 先將數(shù)組的前 k 個元素依次入堆。
3. 從第 k+1 個元素開始，若當(dāng)前元素大于堆頂元素，則將堆頂元素出堆，并將當(dāng)前元素入堆。
4. 遍歷完成后，堆中保存的就是最大的 k 個元素。

圖 8-8   基于堆尋找最大的 k 個元素

總共執(zhí)行了n輪入堆和出堆，堆的最大長度為k，因此時間復(fù)雜度為O(nlogk)。該方法的效率很高，當(dāng)k較小時，時間復(fù)雜度趨向O(n)；當(dāng)k較大時，時間復(fù)雜度不會超過O(nlogn)。

另外，該方法適用于動態(tài)數(shù)據(jù)流的使用場景。在不斷加入數(shù)據(jù)時，我們可以持續(xù)維護(hù)堆內(nèi)的元素，從而實現(xiàn)最大k個元素的動態(tài)更新。

top_k.cpp

/* 基于堆查找數(shù)組中最大的 k 個元素 */
priority_queue<int, vector<int>, greater<int>> topKHeap(vector<int> &nums, int k) {
    priority_queue<int, vector<int>, greater<int>> heap;
    // 將數(shù)組的前 k 個元素入堆
    for (int i = 0; i < k; i++) {
        heap.push(nums[i]);
    }
    // 從第 k+1 個元素開始，保持堆的長度為 k
    for (int i = k; i < nums.size(); i++) {
        // 若當(dāng)前元素大于堆頂元素，則將堆頂元素出堆、當(dāng)前元素入堆
        if (nums[i] > heap.top()) {
            heap.pop();
            heap.push(nums[i]);
        }
    }
    return heap;
}