Problem

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.

Note

Solution

TreeMap

1. Store each nums element and its count in HashMap.

2. Traverse its keySet(), store the count of each key into TreeMap, which means reverse the key-value pairs. And TreeMap will sort the elements by count value.

3. Use pollLastEntry() to get last K entries from TreeMap; then addAll() to put values in ArrayList res.

先按照元素-次数的pair将所有元素存入HashMap，再按照次数-元素pair将哈希表里的所有元素存入TreeMap，然后取TreeMap最后的k个元素返回。

public class Solution {
    public List<Integer> topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        for (int num: nums) {
            map.put(num, map.getOrDefault(num, 0)+1);
        }
        TreeMap<Integer, List<Integer>> sorted = new TreeMap<>();
        for (int num: map.keySet()) {
            int count = map.get(num);
            if (!sorted.containsKey(count)) sorted.put(count, new LinkedList<>());
            sorted.get(count).add(num);
        }
        List<Integer> res = new ArrayList<>();
        while (res.size() < k) {
            Map.Entry<Integer, List<Integer>> entry = sorted.pollLastEntry();
            res.addAll(entry.getValue());
        }
        return res;
    }
}

Min Heap

lambda-expression

public class Solution {
    public List<Integer> topKFrequent(int[] nums, int k) {
        List<Integer> res = new ArrayList<>();
        if (nums.length < k) return res;
        Map<Integer, Integer> map = new HashMap<>();
        for (int num: nums) {
            map.put(num, map.getOrDefault(num, 0)+1);
        }
        PriorityQueue<Map.Entry<Integer, Integer>> minHeap = new PriorityQueue<>((a, b)->(a.getValue()-b.getValue()));
        for (Map.Entry<Integer, Integer> entry: map.entrySet()) {
            minHeap.offer(entry);
            if (minHeap.size() > k) minHeap.poll();
        }
        while (res.size() < k) {
            res.add(minHeap.poll().getKey());
        }
        return res;
    }
}

no-lambda

public class Solution {
    public List<Integer> topKFrequent(int[] nums, int k) {
        List<Integer> res = new ArrayList<>();
        if (nums.length < k) return res;
        Map<Integer, Integer> map = new HashMap<>();
        for (int num: nums) {
            map.put(num, map.getOrDefault(num, 0)+1);
        }
        PriorityQueue<Map.Entry<Integer, Integer>> minHeap = new PriorityQueue<Map.Entry<Integer, Integer>>(new Comparator<Map.Entry<Integer, Integer>>() {
            public int compare(Map.Entry<Integer, Integer> a, Map.Entry<Integer, Integer> b) {
                return a.getValue()-b.getValue();
            }
        });
        for (Map.Entry<Integer, Integer> entry: map.entrySet()) {
            minHeap.offer(entry);
            if (minHeap.size() > k) minHeap.poll();
        }
        while (res.size() < k) {
            Map.Entry<Integer, Integer> entry = minHeap.poll();
            res.add(entry.getKey());
        }
        return res;
    }
}

Max Heap

public class Solution {
    public List<Integer> topKFrequent(int[] nums, int k) {
        List<Integer> res = new ArrayList<>();
        if (nums.length < k) return res;
        Map<Integer, Integer> map = new HashMap<>();
        for (int num: nums) {
            map.put(num, map.getOrDefault(num, 0)+1);
        }
        PriorityQueue<Map.Entry<Integer, Integer>> maxHeap = new PriorityQueue<>((a, b) -> (b.getValue()-a.getValue()));
        for (Map.Entry<Integer, Integer> entry: map.entrySet()) {
            maxHeap.offer(entry);
        }
        while (res.size() < k) {
            res.add(maxHeap.poll().getKey());
        }
        return res;
    }
}