題目網址：https://leetcode.cn/problems/top-k-frequent-elements/

題意：給一 array nums 和一整數 k, 返回出現頻率前 k 高的元素。可以按任何順序返回答案。

進階：請設計 $O(n \cdot log(n))$ time 的演算法, 其中 n 為 nums.size()。

Solution：

想法：利用 Heap

  • 計算每種元素的出現次數
  • 直接將元素 push 到 min heap pq 中
  • 若 pq.size() > k, 代表現在 pq 中有 k + 1 個元素, 故把最小的元素給 pop 掉, 讓 pq 元素個數始終保持在 k 個
  • 此時, pq 中的元素為前 k 大的元素, 將其取出即可
class Solution {
public:
    using pii = pair<int, int>;

    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> freqs; // {num, freq}
        for (const auto& num : nums) {
            ++freqs[num];
        }

        priority_queue<pii, vector<pii>, greater<>> pq; // min heap
        for (const auto& [num, freq] : freqs) {
            pq.emplace(pii{freq, num});
            if (pq.size() > k) { // 讓 pq 始終儲存前 k 大的元素
                pq.pop();
            }
        }

        vector<int> res;
        while (!pq.empty()) {
            res.emplace_back(pq.top().second);
            pq.pop();
        }
        return res;
    }
};
  • time: $O(n \cdot log(k))$ ➔ pq 最多 insert / delete n 個點, 每次操作需 $O(log(k))$, 因為 pq 中最多 k + 1 個點, 其中 n 為 nums 的元素個數
  • space: $O(n)$ ➔ freqs 最大長度為 n, 而 pq 的最大長度為 k + 1, 其中 n ≥ k