Top K element – 陪你刷題

如何辨識 Top K element 類型題目?
題目給你一組資料,要你找出 top / smallest / frequent 的 K 筆資料。這類型題目很好分辨,在 Leetcode 上的統計,這類型題目出現在面試的頻率都相當高,通常有三種解法:

1. 暴力解
2. 使用 Heap
3. Quick Select

Leetcode 邁向千題的關卡,想要把所有題目刷過一遍,對於一位上班族來說就跟寫出來的程式沒有 bug 一樣困難,因此想要將刷題常用的觀念技巧與其對應題目整理出來,一方面可以整理自己思緒,也可以做為未來複習用。
這系列文章會一直持續下去,算是作為工程師職涯的隨身裝備。

如何辨識 Top K element 類型題目

題目給你一組資料,要你找出 top / smallest / frequent 的 K 筆資料。這類型題目很好分辨,在 Leetcode 上的統計,這類型題目出現在面試的頻率都相當高,通常有三種解法:

  1. 暴力解
  2. 使用 Heap
  3. Quick Select

Leetcode #973 K Closest Points to Origin

題目敘述:
平面上有多個點,問你最靠近原點的 K 個點是哪些?

方法 1 Sorting

  • 計算每點與原點的距離,再對所有距離排序,並將前 K 個點作為答案。這是最直覺的解法,但其時間複雜度為 O (N log N) ,空間複雜度 O(1)。
class Solution {
private:
    static bool comp (vector<int> a, vector<int> b)
    {
        return a[0]*a[0] + a[1]*a[1] <
            b[0]*b[0] + b[1]*b[1];
    }
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        std::sort (points.begin(),
            points.end(), comp);
        return vector<vector<int>> (points.begin(),
                                points.begin()+K);
    }
};

方法 2-1 MinHeap

  • 使用 MinHeap 存放所有 element ,再將 heap 的前 K 個元素 pop 出來
  • 時間複雜度為 O (N+ K log N) ,其中 O(N) 為製作 priority queue 的時間,每次由 priority queue 刪除元素要 O (log N) 時間。
  • 空間複雜度 O(N)
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, 
                                int K) {
        priority_queue<vector<int>,
            vector<vector<int>>, heap_compare> 
            pq (points.begin(), points.end());

        vector<vector<int>> result;
        for (int i=0; i<K; i++)
        {
            result.push_back (pq.top());
            pq.pop();
        }
        return result;
    }
private:
    struct heap_compare {
        bool operator () (const vector<int> p1,
                          const vector<int> p2)
        {
            return (p1[0]*p1[0])+(p1[1]*p1[1]) >
                (p2[0]*p2[0])+(p2[1]*p2[1]);
        }
    };
};

方法 2-2 MaxHeap

  • 題目只要找前 K 小的,後面更大元素都不需要保留,使用 MaxHeap 存放 k 個較小的 element ,每新增一個 element 到 heap ,需耗費 O (log K),只要 MaxHeap 的大小超過 K ,就刪除元素,需耗費 O (log K) ,全部時間複雜度為 O (N log K) 。
  • 空間複雜度 O(K) 。

    class Solution {
    public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        priority_queue< pair<vector<int>, int>,
            vector<pair<vector<int>, int>>,
            heap_compare> pq;
    
        for (auto point:points)
        {
            int dis = (point[0]*point[0]) 
                    + (point[1]*point[1]);
            pq.push (make_pair(point, dis));
            if (pq.size() > K)
            {
                pq.pop();
            }
        }
    
        vector<vector<int>> result;
        for (int i=0; i<K; i++)
        {
            result.push_back (pq.top().first);
            pq.pop();
        }
        return result;
    }
    private:
    struct heap_compare {
        bool operator () (const pair<vector<int>, int> p1,
                          const pair<vector<int>, int> p2)
        {
            /*
             *The larger element will be put
             * on the larger index
             */
            return p1.second < p2.second;
        }
    };
    };

方法 3 Quick Select

  • Quick select 一般用於解決 Kth element 的問題, average case 下可以在 O(n) 時間複雜度下完成,worst case 會到 O(n^2) ,可以藉由打亂 array 來盡可能避免 worst case 狀況發生,O(1) 空間複雜度完成。
  • quick select 可以說是 quick sort 的簡化 。 quick sort 會選擇 pivot 並將整個 array 劃為兩邊,接著再以同樣方式處理兩邊資料。但這題要找 Kth element ,並不需要將整個 array 排序,只需找到第 K 大作為 pivot ,將比他小的與比他大的元素分到兩邊即可。
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points,
                                    int K) {
        quickSelect (points, K);
        return vector<vector<int>> (points.begin(),
                                    points.begin() + K);
    }
    bool comp (vector<int> a, vector<int> b)
    {
        return a[0]*a[0] + a[1]*a[1] <
            b[0]*b[0] + b[1]*b[1];
    }

    int partition (vector<vector<int>>& points, 
        int left, int right)
    {
        int shuffle = (rand()%(right-left+1)) + left;
        std::swap (points[right], points[shuffle]);
        int store_index = left;

        for (int i=left; i<right; i++)
        {
            if (comp (points[i], points[right]))
            {
                std::swap (points[i], points[store_index]);
                store_index++;
            }
        }
        std::swap (points[store_index], points[right]);
        return store_index;
    }
    void quickSelect (vector<vector<int>>& points, int k)
    {
        int left = 0, right = points.size()-1;
        while (right > left)
        {
            int pos = partition (points, left, right);
            if (pos == k)
            {
                return;
            }
            else if (pos < k)
            {
                left = pos+1;
            }
            else
            {
                right = pos-1;
            }
        }
    }
};

Leetcode #347 Top K Frequent Elements

方法 1

統計每個 element 出現次數並存到 unordered_map ,請不要對 unordered_map 直接做排序(可參考 stackoverflow 上討論),接著轉存放到 vector 並排序,再挑出前 K 個大的即可。
時間複雜度為 O (N logN) ,空間複雜度為 O (N) 。

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> intDict;
        for (auto num:nums)
        {
            intDict[num]++;
        }
        vector<pair<int,int>> int_arr;
        for (auto intFreq:intDict)
        {
            int_arr.push_back (intFreq);
        }
        sort (int_arr.begin(), int_arr.end(), comp);

        vector<int> result;
        for (int i=0; i<k; i++)
        {
            result.push_back (int_arr[i].first);
        }
        return result;
    }
private:
    static bool comp (pair<int, int> p, pair<int, int> q)
    {
        return p.second > q.second;
    }
};

方法 2 Heap

  • 跟上一題一樣的模板,製造一個 MinHeap ,時間複雜度為 O (N + K log N),空間複雜度為 O (N) 。
class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> intFreq;
        for (auto num:nums)
        {
            intFreq[num]++;
        }

        for (auto freq:intFreq)
        {
            save (freq, k);
        }

        vector<int> ans;
        for (int i=0; i<k; i++)
        {
            ans.push_back (result.top().first);
            result.pop();
        }
        return ans;
    }
private:
    struct heap_comp {
        bool operator() (const pair<int, int> p, const pair<int, int> q) {
            return p.second > q.second;
        }
    };
    priority_queue<pair<int, int>, vector<pair<int, int>>, heap_comp> result;
    void save (const pair<int, int>& input, int k)
    {
        result.push (input);
        if (result.size() > k )
        {
            result.pop ();
        }
    }
};

方法 3 Quick select

時間複雜度可達到 O (N) ,空間複雜度僅 O (1) 。

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map <int, int> intFreq;
        for (auto num:nums)
        {
            intFreq[num]++;
        }
        vector<pair<int, int>> element;
        for (auto freq:intFreq)
        {
            element.push_back
            (make_pair (freq.first, freq.second));
        }
        int kth_smallest_index = element.size() 
                                - k - 1;

        quick_select (0, element.size()-1,
                    kth_smallest_index, element);
        vector<int> result;
        for (int i=0; i<k; i++)
        {
            result.push_back(element[element.size()-1-i].first);
        }
        return result;
    }
private:
    int quick_select (int left, int right, int kth,
                      vector<pair<int, int>>& intDict)
    {
        int pivot;
        while (right > left)
        {
            pivot = partition (left, right, intDict);
            if (pivot == kth)
                break;
            else if (pivot < kth)
            {
                left = pivot+1;
            }
            else
            {
                right = pivot-1;
            }
        }
        return pivot;
    }

    int partition (int left, int right,
                  vector<pair<int, int>>& intDict)
    {
        int shuffle = (rand()%(right-left+1)) + left;
        std::swap (intDict[shuffle], intDict[right]);
        int stored_index = left;
        for (int i=left; i<right; i++)
        {
            if (intDict[i].second <= intDict[right].second)
            {
                std::swap (intDict[i], intDict[stored_index]);
                stored_index++;
            }
        }
        std::swap (intDict[right], intDict[stored_index]);
        return stored_index;
    }
};

方法 4 Bucket sort

桶排序我就不描述,網路上可以找到很多資料

這個方法的時間複雜度可以到 O (n) ,空間複雜度也為 O(n) 。

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> numFreq;
        int max_cnt = 0;
        for (auto num:nums)
        {
            max_cnt = max(max_cnt, ++numFreq[num]);
        }
        vector<vector<int>> bucket(max_cnt+1);
        for (auto num:numFreq)
        {
            bucket[num.second].push_back (num.first);
        }

        vector<int> result;
        for (int i=max_cnt; i>=0 && result.size()<k; i--)
        {
            for (auto num:bucket[i])
            {
                result.push_back (num);
                if (result.size() == k)
                    break;
            }
        }
        return result;
    }
};

延伸練習

Leetcode #703, #692, #215

總結

這類型題目可以使用 sorting, heap 或 quick select,透過練習以上經典的練習題,很快就可以熟悉這三種解法,反倒是要注意 C++ STL 的使用,如果不常使用的話,很容易忘記 STL 的用法,這類型的題目也可以拿來練習如何使用常見的 STL 用法。

Reference

  1. [Java] Three solutions to this classical K-th problem.-Leetcode Discuss
  2. C++ STL, quickselect, priority_queue and multiset
  3. 14 Patterns to Ace Any Coding Interview Question
  4. 花花酱 LeetCode 347. Top K Frequent Elements
  5. 快速排序亲兄弟:快速选择算法

Updated on 2021-07-24 14:40:24 星期六