题目描述
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(nlog(n))
, where
n
<script type="math/tex" id="MathJax-Element-2">n</script> is the array’s size.
算法分析
题目要求我们输出前k个出现最频繁的元素,因为k最大可以等于n,所以在最坏情况下我们需要统计所有数字的出现次数。那么这个问题就分成了两个部分:
- 统计所有不同的数字出现的次数
- 找出出现次数前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;
}
};