347. Top K Frequent Elements

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 is the array's size.

Solution:

The idea is to count the frequency of each number. Then we use a bucket array and put all the numbers that appear x times to index x.

Hence, we add the most frequent number to the result until reach k elements.

Both time complexity and space complexity are O(n).



Code:

public class Solution {
    public List<Integer> topKFrequent(int[] nums, int k) {
        List<Integer> result = new ArrayList<>();
        List<Integer>[] bucket = new List[nums.length + 1];
        HashMap<Integer, Integer> map = new HashMap<>();
        for (int i = 0; i < nums.length; i++) {
            if (!map.containsKey(nums[i])) {
                map.put(nums[i], 1);
            }
            else {
                map.put(nums[i], map.get(nums[i]) + 1);
            }
        }
        
        for (int key : map.keySet()) {
            int frequency = map.get(key);
            if (bucket[frequency] == null) {
                bucket[frequency] = new ArrayList<>();
            }
            bucket[frequency].add(key);
        }
        
        for (int i = bucket.length - 1; i >= 0 && k > 0; i--) {
            if (bucket[i] != null) {
                for (int num : bucket[i]) {
                    result.add(num);
                    k--;
                    if (k == 0) {
                        break;
                    }
                }
            }
        }
        return result;
    }
}