0347. 前 K 个高频元素 Top K Frequent Elements

问题描述

优先级队列

因为要计算 top-k 的问题, 我们自然想到了优先级队列 Priority Queue.

• 同样是先用 hashmap 来统计各整数出现的频率
• 然后将它们存放到一个最大堆 heap 中, 每个元素是一个元组, 包含 (频率, 整数值) 两项, 以频率和整数值的降序来排列
• 之后将最大堆转换成数组, 并截取前 k 个元素即可

Rust 实现

#![allow(unused)]
fn main() {
use std::collections::{BinaryHeap, HashMap};

// HashMap + Priority Queue
// 字典计数
pub fn top_k_frequent1(nums: Vec<i32>, k: i32) -> Vec<i32> {
    assert!(!nums.is_empty());
    assert!(k > 0);

    // 计数
    let mut map: HashMap<i32, usize> = HashMap::new();
    for &num in &nums {
        map.entry(num).and_modify(|count| *count += 1).or_insert(1);
    }

    // 优先队列, BinaryHeap 是一个最大堆
    let k = k as usize;
    let mut heap = BinaryHeap::with_capacity(map.len());
    for (num, count) in map {
        heap.push((count, num));
    }

    // 转换成数组.
    let mut out = Vec::with_capacity(k);
    while let Some(top) = heap.pop() {
        out.push(top.1);
        if out.len() == k {
            break;
        }
    }
    out
}
}

Python 实现

import heapq

class Solution:
    def topKFrequent(self, nums: list[int], k: int) -> list[int]:
        # 首先统计数值的频率
        word_count = {}
        for num in nums:
            word_count[num] = word_count.get(num, 0) + 1
        # 构造最大堆, 堆中的元素是 (频率, 数值)
        pq = [(count, value) for (value, count) in word_count.items()]
        heapq.heapify(pq)
        # 得到最大堆的 top-k
        lst = heapq.nlargest(k, pq)
        # 提取 top-k 中的数值
        return [value for (_count, value) in lst]

C++ 实现

#include <queue>
#include <unordered_map>
#include <vector>

class Solution {
 public:
  std::vector<int> topKFrequent(std::vector<int>& nums, int k) {
    // 先计算各数值出现的频率
    // (number, freq)
    std::unordered_map<int, size_t> freqs;
    for (int num: nums) {
      freqs[num] += 1;
    }

    // 再将它们存入到 priority_queue, 它是个最大堆.
    // (freq, number), 以降序排列
    std::priority_queue<std::pair<size_t, int>> queue;
    for (const auto& pair : freqs) {
      queue.emplace(pair.second, pair.first);
    }

    // 最后导出为 vector
    std::vector<int> out;
    while (!queue.empty() && out.size() < k) {
      out.emplace_back(queue.top().second);
      queue.pop();
    }

    return out;
  }
};

手动对数值频率进行排序

上面的方法使用了最大堆来对数值出现的频率进行了排序, 但我们发现它并不是最快的算法.

• 我们可以将有相同频率的所有数值都存放在同一个数组中
• 然后用一个大的数组来存放, 以数值的频率作为元素的索引

代码实现如下:

#![allow(unused)]
fn main() {
use std::collections::{BinaryHeap, HashMap};

// HashMap + Bucket
pub fn top_k_frequent3(nums: Vec<i32>, k: i32) -> Vec<i32> {
    assert!(!nums.is_empty());
    assert!(k > 0);

    // 计数
    let mut count_map: HashMap<i32, usize> = HashMap::new();
    for &num in &nums {
        *count_map.entry(num).or_insert(0) += 1;
    }

    // 将有相同频率的数值放在一个数组中.
    let max_count: usize = count_map.values().max().copied().unwrap_or_default();
    // 要注意数组的元素个数是 max_count + 1
    let mut count_list = vec![Vec::new(); max_count + 1];
    for (&num, &count) in &count_map {
        count_list[count].push(num);
    }

    // 从最高频率开始, 依次收集整数值.
    let k_usize = k as usize;
    let mut out = Vec::new();
    for array in count_list.into_iter().rev() {
        if !array.is_empty() {
            out.extend(&array);
        }
        if out.len() >= k_usize {
            break;
        }
    }
    out
}
}

Quick Select

TODO

Counting Sort

TODO