算法

(堆) $O(nlog(k))$

用小根堆找出频率最高的前k个元素

Java 代码

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> freq = new HashMap();
        for (int i = 0; i < nums.length; ++i) {
            int f = freq.getOrDefault(nums[i], 0) + 1;
            freq.put(nums[i], f);
        }

        PriorityQueue<Integer> minHeap = new PriorityQueue<>((o1, o2)->Integer.compare(freq.get(o1), freq.get(o2)));

        for (Integer elem : freq.keySet()) {
            minHeap.add(elem);
            if (minHeap.size() > k) {
                minHeap.remove();
            }
        }

        int[] res = new int[minHeap.size()];
        for (int i = 0; i < res.length; ++i) {
            res[i] = minHeap.remove();
        }

        return res;
    }
}

算法2

(STL) $O(n)$

使用 nth_element 使得第 k 大的元素左边所有元素都比它大，并且右边所有元素都比它小，然后截取前k个元素。

C++ 代码

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> freq;
        auto it = nums.begin();
        for (int n : nums) if (!freq[n]++) *it++ = n;
        nums.resize(freq.size());
        nth_element(
            nums.begin(), nums.begin() + (k - 1), nums.end(),
            [&](int a, int b) { return freq[a] > freq[b]; });
        nums.resize(k);
        return move(nums);
    }
};