前 K 个高频元素问题

作者：Grey

原文地址:

博客园：前 K 个高频元素问题

CSDN：前 K 个高频元素问题

题目描述

LeetCode 347. Top K Frequent Elements

思路

第一步，针对数组元素封装一个数据结构

public class Node {
    int v;
    int t;
    public Node(int value, int times) {
        v = value;
        t = times;
    }
}

其中v表示数组元素，t表示数组元素出现的次数。

第二步，使用哈希表把每个元素的词频先存一下。其中key是数组元素，value是Node类型，封装了数组元素和次数。

        Map<Integer, Node> freqMap = new HashMap<>();
        for (int n : arr) {
            if (freqMap.containsKey(n)) {
                // 存在就把词频加一
                freqMap.get(n).t++;
            } else {
                // 不存在就新建一个词频
                freqMap.put(n, new Node(n, 1));
            }
        }

第三步，使用一个小根堆，按词频从小到大。我们需要将这个小根堆维持在K个高频元素。具体做法如下

如果堆未超过K个元素，可以入堆；

如果堆已经到了K个元素了，就看堆顶的元素出现的次数是否比即将要遍历的元素出现的次数少，如果堆顶元素出现的次数比即将要遍历的元素少，

说明即将遍历的元素比堆顶元素更高频，可以替换掉堆顶元素，将其入堆；

如果堆已经超过K个元素了，那么弹出元素，让堆始终保持在K个元素。

第四步，弹出堆中所有元素，即为前K个高频元素。

完整代码如下：

    public static class Node {
        // 值
        public int v;
        // 次数
        public int t;

        public Node(int value, int times) {
            v = value;
            t = times;
        }
    }

    public static int[] topKFrequent(int[] arr, int k) {
        if (arr == null || arr.length == 0 || arr.length < k) {
            return null;
        }
        Map<Integer, Node> freqMap = new HashMap<>();
        for (int n : arr) {
            if (freqMap.containsKey(n)) {
                freqMap.get(n).t++;
            } else {
                freqMap.put(n, new Node(n, 1));
            }
        }
        // 字符种类没有k个，无法得到结果
        if (freqMap.size() < k) {
            return null;
        }
        int[] ans = new int[k];
        PriorityQueue<Node> topK = new PriorityQueue<>(k, Comparator.comparingInt(o -> o.t));
        for (Map.Entry<Integer, Node> entry : freqMap.entrySet()) {
            if (topK.size() <= k || topK.peek().t < entry.getValue().t) {
                topK.offer(entry.getValue());
            }
            if (topK.size() > k) {
                topK.poll();
            }
        }
        int i = 0;
        while (!topK.isEmpty()) {
            ans[i++] = topK.poll().v;
        }
        return ans;
    }