Sign in Subscribe
算法

前 K 个高频元素

loser

• 8 min read
Share

给你一个整数数组 nums 和一个整数 k。

请返回其中出现频率前 k 高的元素。

答案可以按任意顺序返回。

示例 1：

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

示例 2：

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

示例 3：

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

提示：

  • 1 <= nums.length <= 10^5
  • -10^4 <= nums[i] <= 10^4
  • k 的取值范围是 [1, 数组中不相同的元素的个数]
  • 题目数据保证答案唯一，也就是数组中前 k 个高频元素的集合是唯一的

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

解法一（哈希计数 + 小根堆）：

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> cnt;

        for (int num : nums)
        {
            cnt[num]++;
        }

        priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> pq;

        for (auto& [num, freq] : cnt)
        {
            pq.push({freq, num});

            if (pq.size() > k)
            {
                pq.pop();
            }
        }

        vector<int> ans;

        while (!pq.empty())
        {
            ans.push_back(pq.top().second);
            pq.pop();
        }

        return ans;
    }
};

解法二（哈希计数 + 桶排序）：

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> cnt;

        for (int num : nums)
        {
            cnt[num]++;
        }

        vector<vector<int>> buckets(nums.size() + 1);

        for (auto& [num, freq] : cnt)
        {
            buckets[freq].push_back(num);
        }

        vector<int> ans;

        for (int freq = buckets.size() - 1; freq >= 1 && ans.size() < k; freq--)
        {
            for (int num : buckets[freq])
            {
                ans.push_back(num);

                if (ans.size() == k)
                {
                    break;
                }
            }
        }

        return ans;
    }
};

核心思想

这题要求的是出现频率前 k 高的元素，而不是元素本身的大小。

所以第一步一定是统计每个数字出现了多少次。

最直接的做法是：

  1. 用哈希表统计频率。
  2. 把所有不同元素按频率排序。
  3. 取前 k 个。

但如果完整排序，时间复杂度是 O(m log m)，其中 m 是不同元素个数。

最坏情况下 m = n，这就接近 O(n log n)，不满足进阶要求。

这题最关键的观察是：

只需要保留频率最高的 k 个元素，不需要把所有元素完整排序。

所以可以使用大小为 k 的小根堆。

堆里始终保存当前见过的频率最高的 k 个元素。

也可以利用频率不会超过 n 这一点，用桶排序从高频到低频收集答案。

解法一：小根堆

先用哈希表统计频率：

unordered_map<int, int> cnt;

其中：

key   : 数字
value : 出现次数

然后维护一个小根堆：

priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> pq;

堆中元素是：

(频率, 数字)

小根堆的堆顶，是当前堆中频率最低的元素。

每次把一个元素加入堆。

如果堆的大小超过 k，就弹出堆顶：

if (pq.size() > k)
{
    pq.pop();
}

这样堆里最终留下的，就是频率最高的 k 个元素。

为什么小根堆能保留前 k 高频

小根堆中最多保存 k 个元素。

当新元素进入堆后，如果堆大小变成 k + 1，说明当前已经有 k + 1 个候选元素。

其中频率最低的那个不可能属于前 k 高频。

因为至少有另外 k 个元素的频率不低于它。

所以把堆顶弹出是安全的。

重复这个过程后，所有不可能进入前 k 的元素都会被淘汰。

最终堆中剩下的就是前 k 高频元素。

解法二：桶排序

每个元素的出现频率最少是 1，最多是 nums.size()。

所以可以创建 n + 1 个桶：

vector<vector<int>> buckets(nums.size() + 1);

其中：

buckets[freq] 保存所有出现次数为 freq 的数字

统计完频率后，把每个数字放进对应频率的桶里：

buckets[freq].push_back(num);

然后从高频桶往低频桶遍历：

for (int freq = buckets.size() - 1; freq >= 1 && ans.size() < k; freq--)

每遇到一个数字，就加入答案。

直到答案中有 k 个元素为止。

由于频率范围只有 0...n，桶排序可以做到线性时间。

为什么答案可以任意顺序

题目只要求返回前 k 高频元素的集合。

不要求按频率从高到低排列，也不要求按数字大小排列。

所以小根堆解法最后直接把堆中元素取出来即可。

桶排序解法虽然是从高频到低频收集，但同一个频率桶中的元素顺序也不重要。

题目保证答案唯一，意思是前 k 个高频元素的集合唯一。

但输出顺序仍然可以任意。

边界情况

如果 nums 只有一个元素：

nums = [1], k = 1

统计频率后只有一个数字，直接返回 [1]。

如果 k 等于不同元素的个数，那么所有不同元素都应该返回。

如果数组中有负数，不影响哈希表计数。

如果多个元素频率相同，但它们都在前 k 集合内部，返回顺序不重要。

题目保证前 k 高频元素的集合唯一，所以不会出现第 k 名边界无法确定的情况。

正确性证明

我们证明：两个解法都能返回出现频率前 k 高的元素。

结论 1：哈希表能正确统计每个元素的频率

算法遍历 nums 中每个元素 num，执行：

cnt[num]++;

每出现一次，就给对应数字的计数加一。

遍历结束后，cnt[num] 就等于 num 在数组中的出现次数。

因此频率统计正确。

结论 2：小根堆始终保存当前处理过元素中的前 k 高频

处理每个不同元素时，算法把它的 (freq, num) 放入堆中。

如果堆大小没有超过 k，这些元素都暂时保留。

如果堆大小超过 k，算法弹出频率最低的堆顶元素。

这个被弹出的元素在当前 k + 1 个候选中频率最低，不可能属于当前前 k 高频集合。

因此弹出它不会丢掉当前前 k 高频元素。

所以每一步结束后，堆中都保存当前已经处理元素中的前 k 高频。

结论 3：小根堆最终留下的就是全局前 k 高频元素

根据结论 2，当所有不同元素都处理完成后，堆中保存的就是所有不同元素中的前 k 高频。

算法把堆中数字取出作为答案。

因此小根堆解法正确。

结论 4：桶排序从高频到低频收集不会漏掉答案

根据结论 1，每个数字都会被放入自己频率对应的桶中。

桶下标越大，表示频率越高。

算法从最高频率桶开始向低频率桶遍历。

因此先加入答案的元素，频率一定不低于后加入的元素。

当答案数量达到 k 时，已经收集到频率最高的 k 个元素。

所以桶排序解法正确。

得出结论

由结论 1 可知，频率统计正确。

由结论 2 和结论 3 可知，小根堆能正确保留前 k 高频元素。

由结论 4 可知，桶排序也能正确从高到低收集前 k 个元素。

因此两个解法都能正确返回答案。

举例理解

以：

nums = [1,1,1,2,2,3], k = 2

为例。

先统计频率：

1 -> 3
2 -> 2
3 -> 1

前 2 高频元素就是：

1, 2

如果使用小根堆：

加入元素 堆中保留的候选
(3, 1) 1
(2, 2) 1, 2
(1, 3) 加入后超过 k，弹出频率最低的 3，保留 1, 2

最终返回：

[1,2]

复杂度分析

设数组长度为 n，不同元素个数为 m。

解法一

统计频率需要遍历数组一次。

每个不同元素最多入堆一次，堆大小最多为 k。

  • 时间复杂度：O(n + m log k)
  • 空间复杂度：O(m + k)

因为 m <= n，且通常 k 小于 m，该复杂度优于完整排序的 O(n log n)。

解法二

统计频率需要 O(n)。

把不同元素放入桶中需要 O(m)。

从桶中收集答案最多扫描 n 个频率位置，并收集若干元素。

  • 时间复杂度：O(n)
  • 空间复杂度：O(n + m)

桶排序满足进阶要求的线性时间复杂度，是本题最强的复杂度写法。