Leetcode 刷題 pattern - Top K elements


Posted by Po-Jen on 2020-04-12

前言

身在大 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

題目

img

暴力解

這題比較無腦的方法,就是先 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

題目

img

暴力解

這題的暴力法一樣很單純,我們可以用一個 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

題目

img
img

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

題目

img

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 演算法工程師,對機器人、電腦視覺和人工智慧有少許研究,正致力改善 自己的習慣


#Leetcode









Related Posts

[Golang] slice 作為參數傳入 func 的注意事項

[Golang] slice 作為參數傳入 func 的注意事項

我遇過的最難的 Cookie 問題

我遇過的最難的 Cookie 問題

Command Line 超新手入門

Command Line 超新手入門



Newsletter




Comments