Question

347. 前 K 个高频元素

English Version

题目描述

给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。

 

示例 1:

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]

示例 2:

输入: nums = [1], k = 1
输出: [1]

 

提示：

• 1 <= nums.length <= 105
• k 的取值范围是 [1, 数组中不相同的元素的个数]
• 题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的

 

进阶：你所设计算法的时间复杂度 必须 优于 O(n log n) ，其中 n_ _是数组大小。

解法

方法一：哈希表 + 优先队列（小根堆）

使用哈希表统计每个元素出现的次数，然后使用优先队列（小根堆）维护前 $k$ 个出现次数最多的元素。

时间复杂度 $O(n\log k)$。

class Solution:
  def topKFrequent(self, nums: List[int], k: int) -> List[int]:
    cnt = Counter(nums)
    return [v[0] for v in cnt.most_common(k)]

class Solution {
  public int[] topKFrequent(int[] nums, int k) {
    Map<Integer, Long> frequency = Arrays.stream(nums).boxed().collect(
      Collectors.groupingBy(Function.identity(), Collectors.counting()));
    Queue<Map.Entry<Integer, Long>> queue = new PriorityQueue<>(Map.Entry.comparingByValue());
    for (var entry : frequency.entrySet()) {
      queue.offer(entry);
      if (queue.size() > k) {
        queue.poll();
      }
    }
    return queue.stream().mapToInt(Map.Entry::getKey).toArray();
  }
}

using pii = pair<int, int>;

class Solution {
public:
  vector<int> topKFrequent(vector<int>& nums, int k) {
    unordered_map<int, int> cnt;
    for (int v : nums) ++cnt[v];
    priority_queue<pii, vector<pii>, greater<pii>> pq;
    for (auto& [num, freq] : cnt) {
      pq.push({freq, num});
      if (pq.size() > k) {
        pq.pop();
      }
    }
    vector<int> ans(k);
    for (int i = 0; i < k; ++i) {
      ans[i] = pq.top().second;
      pq.pop();
    }
    return ans;
  }
};

func topKFrequent(nums []int, k int) []int {
  cnt := map[int]int{}
  for _, v := range nums {
    cnt[v]++
  }
  h := hp{}
  for v, freq := range cnt {
    heap.Push(&h, pair{v, freq})
    if len(h) > k {
      heap.Pop(&h)
    }
  }
  ans := make([]int, k)
  for i := range ans {
    ans[i] = heap.Pop(&h).(pair).v
  }
  return ans
}

type pair struct{ v, cnt int }
type hp []pair

func (h hp) Len() int       { return len(h) }
func (h hp) Less(i, j int) bool { return h[i].cnt < h[j].cnt }
func (h hp) Swap(i, j int)    { h[i], h[j] = h[j], h[i] }
func (h *hp) Push(v any)    { *h = append(*h, v.(pair)) }
func (h *hp) Pop() any      { a := *h; v := a[len(a)-1]; *h = a[:len(a)-1]; return v }

function topKFrequent(nums: number[], k: number): number[] {
  let hashMap = new Map();
  for (let num of nums) {
    hashMap.set(num, (hashMap.get(num) || 0) + 1);
  }
  let list = [...hashMap];
  list.sort((a, b) => b[1] - a[1]);
  let ans = [];
  for (let i = 0; i < k; i++) {
    ans.push(list[i][0]);
  }
  return ans;
}

use std::collections::HashMap;
impl Solution {
  pub fn top_k_frequent(nums: Vec<i32>, k: i32) -> Vec<i32> {
    let mut map = HashMap::new();
    let mut max_count = 0;
    for &num in nums.iter() {
      let val = map.get(&num).unwrap_or(&0) + 1;
      map.insert(num, val);
      max_count = max_count.max(val);
    }
    let mut k = k as usize;
    let mut res = vec![0; k];
    while k > 0 {
      let mut next = 0;
      for key in map.keys() {
        let val = map[key];
        if val == max_count {
          res[k - 1] = *key;
          k -= 1;
        } else if val < max_count {
          next = next.max(val);
        }
      }
      max_count = next;
    }
    res
  }
}

方法二

class Solution:
  def topKFrequent(self, nums: List[int], k: int) -> List[int]:
    cnt = Counter(nums)
    hp = []
    for num, freq in cnt.items():
      heappush(hp, (freq, num))
      if len(hp) > k:
        heappop(hp)
    return [v[1] for v in hp]

class Solution {
  public int[] topKFrequent(int[] nums, int k) {
    Map<Integer, Integer> cnt = new HashMap<>();
    for (int v : nums) {
      cnt.put(v, cnt.getOrDefault(v, 0) + 1);
    }
    PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[1] - b[1]);
    for (var e : cnt.entrySet()) {
      pq.offer(new int[] {e.getKey(), e.getValue()});
      if (pq.size() > k) {
        pq.poll();
      }
    }
    int[] ans = new int[k];
    for (int i = 0; i < k; ++i) {
      ans[i] = pq.poll()[0];
    }
    return ans;
  }
}

function topKFrequent(nums: number[], k: number): number[] {
  const map = new Map<number, number>();
  let maxCount = 0;
  for (const num of nums) {
    map.set(num, (map.get(num) ?? 0) + 1);
    maxCount = Math.max(maxCount, map.get(num));
  }

  const res = [];
  while (k > 0) {
    for (const key of map.keys()) {
      if (map.get(key) === maxCount) {
        res.push(key);
        k--;
      }
    }
    maxCount--;
  }
  return res;
}