前K个高频元素[小根堆和大根堆的使用]

前言

当题目需要有序性时，果断排序，可二分快速寻找答案，或是利用有序性快速判定，如双指针等等。除此之外，小根堆和大根堆是取K个的利器，原型基于树状数组–堆结构。

一、前K个高频元素

二、小根堆&大根堆

1、O(KlogN)大根堆

// 前K个高频元素
// 当需要有序性的时候，果断排序。
public class TopKFrequent {
    // O(KlogN)，让大顶堆不断出K个元素。
    public int[] topKFrequent(int[] nums, int k) {
        // 先预处理数据。
        Map<Integer, Integer> m = new HashMap<>();
        for (int num : nums) m.put(num, m.getOrDefault(num, 0) + 1);
        // 把所有数据装入大顶堆。
        Set<Map.Entry<Integer, Integer>> entries = m.entrySet();
        PriorityQueue<int[]> que = new PriorityQueue<>((o1, o2) -> o2[1] - o1[1]);
        for (Map.Entry<Integer, Integer> entry : entries) {
            que.offer(new int[]{entry.getKey(), entry.getValue()});
        }

        int[] r = new int[k];
        for (int i = 0; i < k; i++) {
            r[i] = que.poll()[0];
        }
        return r;
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

2、O(NlogK)小根堆

// O(NlogK)，让小顶堆的元素不超过K个。
class TopKFrequent2 {
    public int[] topKFrequent(int[] nums, int k) {
        // 先预处理数据。
        Map<Integer, Integer> m = new HashMap<>();
        for (int num : nums) m.put(num, m.getOrDefault(num, 0) + 1);
        // 把数据装入小顶堆。
        Set<Map.Entry<Integer, Integer>> entries = m.entrySet();
        PriorityQueue<int[]> que = new PriorityQueue<>(Comparator.comparingInt(o -> o[1]));
        for (Map.Entry<Integer, Integer> entry : entries) {
            // 让堆中的元素不超过K个。
            if (que.size() == k && que.peek()[1] < entry.getValue()) que.poll();
            if (que.size() != k) que.offer(new int[]{entry.getKey(), entry.getValue()});
        }

        //遍历小顶堆所有元素
        int[] r = new int[k];
        for (int i = 0; i < k; i++) r[i] = que.poll()[0];

        return r;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

总结

1）当需要有序性的时候，果断排序。

2）小根堆和大根堆的基本使用场景和方式，Java中又现成的API，即优先队列PriorityQueue类。

参考文献

[1] LeetCode 前K个高频元素