前K个最频繁的元素

题目描述

Given a non-empty array of integers, return the k most frequent elements.

For example,
Given [1,1,1,2,2,3] and k = 2, return [1,2].

Note:
* You may assume k is always valid, 1 ≤ k ≤ number of unique elements.
* Your algorithm’s time complexity must be better than $O(n\log (n))$, where n <script type="math/tex" id="MathJax-Element-2">n</script> is the array’s size.

算法分析

题目要求我们输出前k个出现最频繁的元素，因为k最大可以等于n，所以在最坏情况下我们需要统计所有数字的出现次数。那么这个问题就分成了两个部分：

1. 统计所有不同的数字出现的次数
2. 找出出现次数前k大的数字

对于问题1，因为数字可能很大，我们需要借助HashMap进行统计，时间复杂度是O(n)的。对于问题2，有多种方法：一种简单的方法是，对所有的次数快速排序，然后输出前k个，这样的时间复杂度是O(n log n)，不符合本题的要求。我们需要进行优化。

因为最后只需要返回k个数字，所以我们只需要一直维护一个大小为k的小根堆。当新的数字出现的次数大于堆中最小的次数时，我们对堆进行更新。时间复杂度是O(n log k)，是符合题目要求的。

那有没有办法进一步优化呢？因为k最坏情况下还是等于n的，n log k不是很理想。那么我们就需要换一种排序的方法。有一种排序的方法，其复杂度只和需要排序的数字的大小有关，而在本题中，需要排序的数字大小至多为n（某个数出现了n次）。答案是桶排序！桶排序就是用一个数组bucket记录每个数字出现的次数，每次把数字丢到相应编号的桶中，然后从后往前穷举每一个桶，取出其中的元素直到取满k个。时间复杂度是O(n)。

最后本问题的最优算法的时间复杂度是O(n)。

参考程序

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        vector<int> res;
        if (!nums.size()) return res;
        unordered_map<int, int> cnt;
        for (auto num : nums) cnt[num]++;
        vector<vector<int>> bucket(nums.size() + 1);
        for (auto kv : cnt) {
            bucket[kv.second].push_back(kv.first);
        }

        for (int i = bucket.size() - 1; i >= 0; --i) {
            for (int j = 0; j < bucket[i].size(); ++j){
                res.push_back(bucket[i][j]);
                if (res.size() == k) return res;
            }
        }

        return res;
    }
};

题目来源

LeetCode 347. Top K Frequent Elements