前 K 个高频元素

题目要求

你需要编写一个函数,该函数接收两个参数:一个整数数组 nums 和一个整数 k。你的任务是找出数组 nums 中出现频率最高的 k 个元素,并返回这些元素。结果的输出顺序不做要求。

解题思路

要解决这个问题,我们可以采用以下步骤:

  1. 统计元素频率:首先,我们需要遍历数组 nums,使用一个哈希表(或字典)来统计每个元素出现的次数。

  2. 构建频率数组:然后,我们可以根据元素的频率,将它们放入一个“频率数组”中,其中索引对应于频率,值是具有该频率的元素列表。

  3. 使用堆(优先队列):为了不遍历整个频率数组,我们可以使用一个最小堆(或优先队列)来维护当前最高的 k 个频率的元素。堆中的每个元素是一个频率和元素值的对。

  4. 遍历哈希表:遍历哈希表,对于每个元素,我们将其频率和它本身作为一个对插入到最小堆中。如果堆的大小超过了 k,我们就弹出堆顶元素(即频率最小的元素)。

  5. 提取结果:最后,当我们遍历完哈希表后,堆中剩下的就是频率最高的 k 个元素。我们只需要将这些元素取出并返回即可。

这种方法的时间复杂度主要在于哈希表的构建(O(n)),以及最小堆的维护(O(n log k)),因此总的时间复杂度为 O(n log k)。这是一个相对高效的解法,特别是当 k 远小于 n 的时候。

Golang 版本

package main

import (
	"container/heap"
	"fmt"
)

// An IntHeap is a min-heap of ints.
type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// Function to return the k most frequent elements
func topKFrequent(nums []int, k int) []int {
	// Frequency map
	freqMap := make(map[int]int)
	for _, num := range nums {
		freqMap[num]++
	}

	// Initialize min heap
	h := &IntHeap{}
	heap.Init(h)

	// Map to keep track of frequency to elements
	freqToNums := make(map[int][]int)
	for num, freq := range freqMap {
		freqToNums[freq] = append(freqToNums[freq], num)
		if h.Len() < k {
			heap.Push(h, freq)
		} else if (*h)[0] < freq {
			heap.Push(h, freq)
			heap.Pop(h)
		}
	}

	// Extract the top k frequent elements
	var topK []int
	for h.Len() > 0 {
		freq := heap.Pop(h).(int)
		topK = append(topK, freqToNums[freq]...)
	}

	return topK
}

func main() {
	nums := []int{1, 1, 1, 2, 2, 3}
	k := 2
	fmt.Println(topKFrequent(nums, k)) // Output: [1 2]
}

这段代码定义了一个最小堆,并使用 Go 语言的标准库 container/heap 来实现。topKFrequent 函数首先统计每个数字的频率,然后使用最小堆来找出频率最高的 k 个数字。最后,它返回这些数字。在 main 函数中,我们调用 topKFrequent 函数并打印结果。

C++版本

#include <iostream>
#include <vector>
#include <unordered_map>
#include <queue>
#include <functional>

using namespace std;

vector<int> topKFrequent(vector<int>& nums, int k) {
    // Build the frequency map
    unordered_map<int, int> freqMap;
    for (int num : nums) {
        freqMap[num]++;
    }

    // Comparator for the min heap
    auto comp = [&freqMap](int n1, int n2) { return freqMap[n1] > freqMap[n2]; };
    priority_queue<int, vector<int>, decltype(comp)> minHeap(comp);

    // Build the min heap with the most frequent elements
    for (auto& [num, freq] : freqMap) {
        minHeap.push(num);
        if (minHeap.size() > k) {
            minHeap.pop();
        }
    }

    // Extract the top k elements from the min heap
    vector<int> topK;
    while (!minHeap.empty()) {
        topK.push_back(minHeap.top());
        minHeap.pop();
    }

    return topK;
}

int main() {
    vector<int> nums = {1, 1, 1, 2, 2, 3};
    int k = 2;
    vector<int> topK = topKFrequent(nums, k);
    for (int num : topK) {
        cout << num << " ";
    }
    cout << endl;
    return 0;
}

这段代码首先使用一个 unordered_map 来统计每个数字出现的频率。然后,它定义了一个最小堆(使用 priority_queue),并提供了一个自定义的比较器,该比较器基于频率来比较元素。接着,它将元素插入到最小堆中,保持堆的大小为 k。最后,它从最小堆中提取出前 k 个最频繁的元素,并将它们存储在一个向量中返回。在 main 函数中,我们调用 topKFrequent 函数并打印结果。

Python 版本

from collections import Counter
import heapq

def topKFrequent(nums, k):
    # Count the frequency of each element
    freq_map = Counter(nums)

    # Use a heap to store the frequency and corresponding element
    # Since heapq in Python is a min-heap, we use negative frequencies
    heap = [(-freq, num) for num, freq in freq_map.items()]
    heapq.heapify(heap)

    # Extract the top k elements with the highest frequency
    top_k = [heapq.heappop(heap)[1] for _ in range(k)]

    return top_k

# Example usage:
nums = [1, 1, 1, 2, 2, 3]
k = 2
print(topKFrequent(nums, k))  # Output: [1, 2]

这段代码使用 Python 的 collections.Counter 类来统计每个元素的频率。然后,它创建了一个由元素和其负频率组成的最小堆。由于 Python 的 heapq 模块实现的是最小堆,为了得到频率最高的元素,我们将频率取负。最后,它从堆中弹出前 k 个元素,这些元素就是频率最高的 k 个元素。在示例用法中,我们调用 topKFrequent 函数并打印结果。

Java 版本

import java.util.*;

public class Solution {
    public List<Integer> topKFrequent(int[] nums, int k) {
        // Count the frequency for each element
        Map<Integer, Integer> frequencyMap = new HashMap<>();
        for (int num : nums) {
            frequencyMap.put(num, frequencyMap.getOrDefault(num, 0) + 1);
        }

        // Create a min heap with a comparator based on frequency
        PriorityQueue<Integer> heap = new PriorityQueue<>(Comparator.comparingInt(frequencyMap::get));

        // Keep k top frequent elements in the heap
        for (int num : frequencyMap.keySet()) {
            heap.add(num);
            if (heap.size() > k) {
                heap.poll();
            }
        }

        // Extract the top k frequent elements
        List<Integer> topK = new ArrayList<>();
        while (!heap.isEmpty()) {
            topK.add(heap.poll());
        }
        Collections.reverse(topK); // The smallest frequency is at the front, so reverse to get correct order
        return topK;
    }

    public static void main(String[] args) {
        Solution solution = new Solution();
        int[] nums = {1, 1, 1, 2, 2, 3};
        int k = 2;
        List<Integer> topK = solution.topKFrequent(nums, k);
        for (int num : topK) {
            System.out.print(num + " ");
        }
    }
}

这段代码首先使用一个 HashMap 来统计每个元素的频率。然后,它创建了一个最小堆(使用 PriorityQueue),并提供了一个基于频率的比较器。接着,它将元素插入到最小堆中,保持堆的大小为 k。最后,它从最小堆中提取出前 k 个最频繁的元素,并将它们存储在一个列表中返回。在 main 方法中,我们创建了 Solution 类的实例,调用 topKFrequent 方法并打印结果。

Rust 版本

use std::collections::{BinaryHeap, HashMap};
use std::cmp::Reverse;

fn top_k_frequent(nums: Vec<i32>, k: i32) -> Vec<i32> {
    let mut freq_map = HashMap::new();

    // Count the frequency of each element
    for num in nums {
        *freq_map.entry(num).or_insert(0) += 1;
    }

    // By default Rust provides a max heap, so we use Reverse to turn it into a min heap
    let mut heap = BinaryHeap::new();

    // Keep only the k most frequent elements in the heap
    for (num, &freq) in freq_map.iter() {
        heap.push(Reverse((freq, *num)));
        if heap.len() as i32 > k {
            heap.pop();
        }
    }

    // Extract the elements from the heap
    heap.into_iter().map(|Reverse((_, num))| num).collect()
}

fn main() {
    let nums = vec![1, 1, 1, 2, 2, 3];
    let k = 2;
    let top_k = top_k_frequent(nums, k);
    println!("{:?}", top_k); // Output should be [1, 2] or [2, 1]
}

这段代码首先使用 Rust 的 HashMap 来统计每个元素的频率。然后,它创建了一个 BinaryHeap,Rust 的 BinaryHeap 是一个最大堆,为了实现最小堆的效果,我们使用 Reverse 结构体来反转比较的顺序。接着,它将元素以 (频率, 元素) 的形式插入到堆中,并保持堆的大小为 k。最后,它从堆中提取出所有元素,并只保留元素部分,形成一个向量返回。在 main 函数中,我们调用 top_k_frequent 函数并打印结果。

总结

上述解法的核心思想是使用一个哈希表来统计每个元素在数组中出现的频率,然后利用一个堆(优先队列)来维护出现频率最高的 k 个元素。

  1. 哈希表统计频率:遍历数组,使用哈希表记录每个元素出现的次数。

  2. 构建堆:根据元素的频率构建一个最小堆(在 Rust 中使用最大堆的 Reverse 特性来实现)。堆中存储的是一个包含频率和元素值的元组(或在其他语言中是对象或结构体)。

  3. 维护堆的大小:在将元素添加到堆中时,一旦堆的大小超过 k,就移除堆顶元素,这样可以确保堆中始终保留频率最高的 k 个元素。

  4. 提取结果:最后,从堆中提取所有元素,这些就是出现频率最高的 k 个元素。在 Rust 中,由于使用了最大堆的 Reverse 特性,提取时不需要额外操作;而在其他语言中,可能需要将最小堆中的元素反转以得到正确的顺序。

这种解法的时间复杂度通常是 O(n log k),其中 n 是数组的长度,因为每次堆操作(插入或删除)的时间复杂度是 O(log k)。这种方法在 k 远小于 n 的情况下特别有效,因为它不需要对整个数组进行排序。