Leetcode 刷題 pattern - Top K elements

前言

身在大 CS 時代，有越來越多人投入刷題的行列，在眼花撩亂的題海中，要想有效率地刷題，就需要掌握題目解法中，可以在許多地方應用的觀念，才能以簡禦繁。 Huli 寫的 程式解題新手入門注意事項 和 當我們在學程式時，要學的到底是什麼？ 講得很好，寫題目是為了學會解題的思考法，確保自己掌握重要的資料結構跟演算法。這也是我寫這系列文章的原因，我想把多個散落在各處的題目銜接起來，以後看到相似的問題就可以舉一反三，而不是去背各題目的解法。

另外，這系列文章提供的做法雖然可以讓面試準備事半功倍，但不一定適合你。如果你比較喜歡透過刷很多題來鍛鍊演算法跟資料結構的直覺，那你可以參考 這篇短文，踏上刷遍一千題、把 computational thinking 刻進腦袋中的道路。

接下來，就讓我們進入今天的正題 - Top K elements。

Top K elements 簡介

Top K elements 這個 pattern 會用到 heap，如果你不知道什麼是 heap，那請你先去學習一下，這篇文章 還不錯可以看看。

Top K element 是指題目中通常會要求我們取第 k 大，或者前 k 個 object（根據某種題目給定的排序法），所以我們可以用 heap 自動把 object 排好序，就可以很方便地得到答案，而不需要每次都重新 sort array。

這個 pattern 的概念很簡單直觀，只要你夠熟，只要看到這類的題目一定會聯想到 heap，比較會有變化的地方，就是要怎麽對各種類型的 object 排序，所以你必須要對自定義 heap 排序的實作方法很熟悉。接下來我們直接透過幾個範例來學習。

Top K element 的第一個範例 - Leetcode #215 - Kth Largest Element in an Array

題目

暴力解

這題比較無腦的方法，就是先 sort array，然後直接從裡面取第 k 大的：

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());
        return nums[nums.size()-k];
    }
};

這樣做的時間複雜度是 O(nlogn)，空間複雜度是 O(1)，接下來讓我們看看 Top K element 的 heap 解法。

Top K element 解法

概念上，我們可以儲存一個只包含前 k 大 element 的 min heap，所以這個 heap 的頂端就是我們要的答案。

要建立這個 heap 也很簡單，我們先建立一個 min heap，然後當我們走過 array，就依序把 element 放進這個 heap，只要 heap 的大小超過 k，就把多出來的那個 element 丟掉，直到走完 array，我們就獲得了理想中的 min heap。最後回傳 heap 的頂端就好。

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<int>> minHeap;

        for(auto& n: nums) {
            minHeap.push(n);

            if(minHeap.size() > k) {
                minHeap.pop();
            }
        }

        return minHeap.top();
    }
};

這樣做的時間複雜度是 O(nlogk)，空間複雜度是 O(k)，乍看之下，好像沒有比無腦 sort 的 O(nlogn) 好很多，可是如果 n 比 k 大很多很多，例如 n 是一百萬，k 是三，如果我用 sort 的方法，我得先把一百萬個 element 都先 sort 完，再找第三大的；反之如果用 heap，我隨時只需要維持 heap 裡面的排序就好，因為 heap 只有三個 element，所以維持這個排序所需要的 O(logk) 時間相對就很少，用 heap 的好處就顯現出來了！

這題其實還有更好的解法，如果有興趣可以去看這個 discussion post。

Top K element 的第二個範例 - Leetcode #703 - Kth Largest Element in a Stream

題目

暴力解

這題的暴力法一樣很單純，我們可以用一個 array 來存 element，並 sort 好，如果有超過 k 個 element，就把 array 中第 k+1 小以後的 element 都丟掉。之後加入的新 element 只要跟 array 中第 k 大的比就可以決定要不要加入，一旦加入新的 element 就必須重新 sort。

Top K element 解法

Top K element 的模板解就是一樣用 heap，永遠保持前 k 大的 element 在 heap 中，實作上跟上一題幾乎一樣，不過兩題一起寫，熟悉度就能加倍：

class KthLargest {
public:
    KthLargest(int k, vector<int>& nums) {
        K = k;

        for(int i = 0; i < nums.size(); ++i) {
            minHeap.push(nums[i]);
            if(minHeap.size() > K) {
                minHeap.pop();
            }
        }
    }

    int add(int val) {
        minHeap.push(val);
        if(minHeap.size() > K) {
            minHeap.pop();
        }

        return minHeap.top();
    }

private:
    int K;
    priority_queue<int, vector<int>, greater<int>> minHeap; 
};

Top K element 的第三個範例 - Leetcode #973 - K Closest Points to Origin

題目

Top K element 解法

這題就不多說暴力解啦，總之就是根據距離來排序。

如果要用 heap 的話，我們可以把 <點, 到原點的距離> 放進 max heap 中，裡面存的就是距離原點前 K 近的點。因為我們要比較的是 <點, 到原點的距離> 這種特殊的資料結構，所以 heap 的 comparator 要自己實作。這一題是很好的練習，因為 heap 在很多地方都用得到，所以一定要會自己實作比較的方法。下面的實作是一個參考範例：

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        // Store <point, dist> pair into maxHeap
        priority_queue< pair<vector<int>, float>, vector<pair<vector<int>, float>>, comp > maxHeap;

        for(auto& p: points) {
            float dist = sqrt( pow(p[0],2) + pow(p[1], 2) );
            maxHeap.push(make_pair(p, dist));

            if(maxHeap.size() > K) {
                maxHeap.pop();
            }
        }

        // Put all points in maxHeap into the result
        vector<vector<int>> res;
        for(int i = 0; i < K; i++) {
            res.push_back(maxHeap.top().first);
            maxHeap.pop();
        }

        return res;
    }

private:
    struct comp {
        bool operator() (const pair<vector<int>, float>& p1, 
                         const pair<vector<int>, float>& p2) {
            return p1.second < p2.second;
        }
    };
};

Top K element 的第四個範例 - Leetcode #347 - Top K Frequent Elements

題目

Top K element 解法

我們可以把 <num, freq> 放進 min heap 中，裡面存的就是 frequency 前 K 大的 num。這個 heap 的 comparator 我們一樣要自己實作。

因為一開始還得先算出每個 num 的 freq，所以需要先花 O(n) 的時間建立起 num, freq 的 pair，時間複雜度是 O(n + n*logk)。

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        // Get the freq of each num
        unordered_map<int, int> freqMap;
        for(auto& n: nums) {
            freqMap[n]++;
        }

        // Put all <num, freq> pair into min heap
        priority_queue<pair<int, int>, vector<pair<int, int>>, comp> minHeap;
        for(auto& p: freqMap) {
            minHeap.push(p);

            if(minHeap.size() > k) {
                minHeap.pop();
            }
        }

        // Put all num in min heap into result
        vector<int> res;
        for(int i = 0; i < k; ++i) {
            res.push_back(minHeap.top().first);
            minHeap.pop();
        }

        return res;
    }

private:
    struct comp {
        bool operator() (const pair<int, int>& p1, const pair<int, int>& p2) {
            return p1.second > p2.second;
        }  
    };
};

總結

今天跟大家介紹了一個新的 pattern - Top K element，順便提供幾個經典的可應用題，讓大家體驗一下用這個 pattern 的好處。希望大家在刷題的時候也能多掌握解題的 pattern，用一個 pattern 打掉多題，體會到隱藏在這些 pattern 後的演算法和資料結構的優美之處。

延伸閱讀

1. Leetcode 刷題 pattern - Two Pointer
2. Leetcode 刷題 pattern - Sliding Window
3. Leetcode 刷題 pattern - Next Greater Element
4. Leetcode 刷題 pattern - Fast & Slow Pointer
5. Leetcode 刷題 pattern - Breadth-First Search
6. Leetcode 刷題 pattern - Merge Intervals
7. Leetcode 刷題 pattern - Cyclic Sort
8. Leetcode 刷題 pattern - Two Heaps

關於作者：
@pojenlai 演算法工程師，對機器人、電腦視覺和人工智慧有少許研究，正致力改善 自己的習慣。