347. Top K Frequent Elements

Leetcode HashTable Heap

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]

Example 2:

Input: nums = [1], k = 1
Output: [1]

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.

分析

这道题目常见的做法是利用二叉堆维持频数最大的$k$个元素。为了获取各个元素的频数，需要利用哈希表统计并保存各元素的出现次数。计算频数的时间复杂度是$O(n)$，操作二叉堆的时间复杂度是$O(n\log k)$.

public List<Integer> topKFrequent(int[] nums, int k) {
    List<Integer> list = new ArrayList<>();
    if (nums == null || nums.length == 0 || k <= 0) return list;

    // 计算频数
    HashMap<Integer, Integer> map = new HashMap<>();
    for (int num : nums)
        map.put(num, map.getOrDefault(num, 0) + 1);

    // 用二叉堆维持频数最大的K个元素
    PriorityQueue<int[]> pq = new PriorityQueue<>((o1,o2)->(o1[1]-o2[1]));
    for (int num : map.keySet()) {
        pq.offer(new int[]{num, map.get(num)});
        if (pq.size() > k) pq.poll();
    }

    // 频数最大的K个元素
    while (!pq.isEmpty()) list.add(pq.poll()[0]);
    return list;
}

另一个更好的办法是用桶排序。利用桶排序算法将频数排序，时间和空间复杂度是$O(n)$。

public List<Integer> topKFrequent(int[] nums, int k) {
    List<Integer> list = new ArrayList<>();
    if (nums == null || nums.length == 0 || k <= 0) return list;

    // 计算频数
    HashMap<Integer, Integer> map = new HashMap<>();
    for (int num : nums)
        map.put(num, map.getOrDefault(num, 0) + 1);

    // 桶排序
    List<List<Integer>> buckets = new ArrayList<>();
    for (int i = 0; i <= nums.length; i++) buckets.add(new ArrayList<>());
    for (Map.Entry<Integer, Integer> entry : map.entrySet())
        buckets.get(entry.getValue()).add(entry.getKey());
    for (int i = nums.length; i >= 0 && list.size() < k; i--)
        if (!buckets.get(i).isEmpty()) list.addAll(buckets.get(i));

    //如果数量不唯一 return list.sublist
    return list;
}