Leetcode 347. 前 K 个高频元素

本题链接

题目

给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。

示例 1:

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

示例 2:

输入: nums = [1], k = 1
输出: [1]

提示：

1 <= nums.length <= 105
k 的取值范围是 [1, 数组中不相同的元素的个数]
题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的

进阶：你所设计算法的时间复杂度 必须 优于 O(n log n) ，其中 n 是数组大小。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/top-k-frequent-elements
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

---

思路

首先用一个 unordered_map 来统计数组中各个整数的出现次数。这样问题就转化成了：取出一个 unordered_map 中 value 最大的 K 个 key

其实就是 Top K 问题，三种思路：

• 把所有键值对保存到数组中，直接排序
• 遍历所有键值对，维护一个大小为 K 的小顶堆，遍历结束后小顶堆中存放的就是结果
• 快速选择算法：Leetcode 215. 数组中的第K个最大元素

C++ 代码

直接排序

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        vector<int> ans;
        unordered_map<int, int> numToCnt;
        for (const int& num : nums)
            ++numToCnt[num];
        vector<vector<int>> numAndCnts;
        for (auto it = numToCnt.cbegin(); it != numToCnt.cend(); ++it) {
            numAndCnts.emplace_back(vector<int>{it->first, it->second});
        }
        sort(numAndCnts.begin(), numAndCnts.end(),
            [](const vector<int>& a, const vector<int>& b) -> bool {
                return a[1] > b[1]; });
        for (int i = 0; i < k; ++i)
            ans.emplace_back(numAndCnts[i][0]);
        return ans;
    }
};

维护小顶堆

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        vector<int> ans;
        unordered_map<int, int> numToCnt;
        for (const int& num : nums)
            ++numToCnt[num];
        auto cmp = [](const pair<int, int>& a, const pair<int, int>& b) {
            return a.second > b.second; };
        priority_queue<pair<int, int>, vector<pair<int, int>>,
            decltype(cmp)> q(cmp);
        for (auto it = numToCnt.cbegin(); it != numToCnt.cend(); ++it) {
            if (q.size() < k) {
                q.push(*it);
            } else if (it->second > q.top().second) {
                q.pop();
                q.push(*it);
            }
        }
        while (!q.empty()) {
            ans.emplace_back(q.top().first);
            q.pop();
        }
        return ans;
    }
};

快速选择

class Solution {
public:
    int partition(vector<pair<int, int>>& vec, int l, int r) {
        swap(vec[l], vec[l + rand() % (r - l + 1)]);
        pair<int, int> pivot = vec[l];
        while (l < r) {
            while (l < r && vec[r].second <= pivot.second)
                --r;
            vec[l] = vec[r];
            while (l < r && vec[l].second > pivot.second)
                ++l;
            vec[r] = vec[l];
        }
        vec[l] = pivot;
        return l;
    }

    vector<int> topKFrequent(vector<int>& nums, int k) {
        srand(time(nullptr));
        vector<int> ans;
        unordered_map<int, int> numToCnt;
        for (const int& num : nums)
            ++numToCnt[num];
        vector<pair<int, int>> numAndCnts;
        for (const auto& item : numToCnt)
            numAndCnts.emplace_back(item);
        int lo = 0, hi = numAndCnts.size() - 1;
        int mid = partition(numAndCnts, lo, hi);
        while (mid != k - 1) {
            if (mid < k - 1)
                lo = mid + 1;
            else
                hi = mid - 1;
            mid = partition(numAndCnts, lo, hi);
        }
        for (int i = 0; i <= mid; ++i)
            ans.emplace_back(numAndCnts[i].first);
        return ans;
    }
};