Top K Frequent Elements

Question

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

Example 1:
Input: nums = [1,1,1,2,2,3], k = 2
Output: [1,2]

LeetCodeleetcode.com/problems/top-k-frequent-elements

Solution

/*
@param: int[] 
@return: int[]
Algorithm: 经典的Kth题目，利用heap来不断更新k elements。这里有个小技巧是可以用map.get来比较排序
*/
public int[] topKFrequent(int[] nums, int k) {
    if (nums.length == 0 || k > nums.length) {
        return new int[] {};
    }
    
    Map<Integer, Integer> map = new HashMap<>();
    PriorityQueue<Integer> heap = new PriorityQueue<>((o1,o2) -> map.get(o1)-map.get(o2));
    
    for (int i = 0; i < nums.length; i++) {            
        map.put(nums[i], map.getOrDefault(nums[i], 0)+1);            
    }
    
    for (int n : map.keySet()) {
        heap.offer(n);
        if (heap.size() > k) {
            heap.poll();
        }
    }
    
    int[] res = new int[heap.size()];
    int i = 0;
    for (int n : heap) {
        res[i++] = n;
    }
    return res;
}