【LeetCode-347】前 K 个高频元素

问题

给定一个非空的整数数组，返回其中出现频率前 k 高的元素。

示例

输入： nums = [1,1,1,2,2,3], k = 2
输出： [1,2]

解答

class Solution {
public:
    typedef pair<int, int> P; // 少写点字
    struct cmp {bool operator()(P a,P b) {return a.second > b.second;}}; // 自定义优先队列排序
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> ump;
        priority_queue<P, vector<P>, cmp> pq;
        vector<int> res;
        for (int i : nums) ump[i]++;
        for (auto [f, s] : ump) {
            if (pq.size() < k || pq.top().second < s) pq.push(make_pair(f, s));
            if (pq.size() > k) pq.pop();
        }
        while (!pq.empty()) {
            res.push_back(pq.top().first);
            pq.pop();
        }
        return res;
    }
};

重点思路

典型的TopK问题，可以使用最大堆最小堆来解决。本文要求的是频数的最大K，那么我们可以把数字与其出现的次数构成一个pair，求频数的最大值，然后通过pair输出频数对应的数字。这里有一个知识点是优先队列的自定义排序。需要注意的是这里不能像sort函数之类的直接用lambda函数，优先队列默认为最大堆，可以看做队列越往后越大，自定义排序可以定义为a < b；如果要使用最小堆，越往后越小，则定义为a > b。