leetcode347：前K个高频元素

题目链接

leetcode

题目描述

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

C++ 代码

#include <iostream>
#include <vector>
#include <queue>
#include <unordered_map>
#include <algorithm>
using namespace std;

void printArray(vector<int>& nums) {
    for (const int& num : nums) {
        cout << num << " ";
    }
    cout << endl;
}

/*
    哈希表 + 堆
    堆中保存的结构体需要自定义（定义成员、定义比较规则）
    
    时间复杂度：O(nlogk)
        n 为数组的长度，k 为高频元素的个数。
    空间复杂度：O(n)
        哈希表的大小为 O(n)，堆的大小为 O(k)，共计为 O(n)。
*/
class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        for (const int& num : nums) {
            num2cnt[num] ++;
        }

        for (const auto& it : num2cnt) {
            heap.push({it.second, it.first});
            if (heap.size() > k) {
                heap.pop();
            }
        }

        vector<int> res;
        res.reserve(k);
        while (heap.size()) {
            Node node = heap.top(); heap.pop();
            res.push_back(node.num);
        }
        // reverse(res.begin(), res.end()); // 升序排列

        return res;
    }

private:
    struct Node {
        int cnt;
        int num;
        
        bool operator<(const Node& t) const {
            if (cnt != t.cnt) return cnt > t.cnt;
            return num < t.num;
        }
    };

    priority_queue<Node> heap;       // 大根堆
    unordered_map<int, int> num2cnt; // 哈希表
};

/*
    思路同上
    堆中保存元素类型是 pair<int, int>，自定义 priority_queue 排序规则
*/
class Solution_1 {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> num2cnt;
        for (const int& num : nums) {
            num2cnt[num] ++;
        }

        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(&cmp)> heap(cmp);
        for (const auto& it : num2cnt) {
            int num = it.first, cnt = it.second;
            heap.push({num, cnt});
            if (heap.size() > k) {
                heap.pop();
            }
        }

        vector<int> res;
        res.reserve(k);
        while (heap.size()) {
            res.emplace_back(heap.top().first);
            heap.pop();
        }
        // reverse(res.begin(), res.end()); // 升序排列

        return res;
    }

private:
    static bool cmp(pair<int, int>& a, pair<int, int>& b) {
        return a.second > b.second;
    }
};

/*
    思路同上，结果数组不排序
*/
class Solution_2 {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> num2cnt;
        for (const int& num : nums) {
            num2cnt[num] ++;
        }

        // 自定义 priority_queue 排序规则
        auto comp = [](const pair<int, int>& a, const pair<int, int>& b) {
            return a.second > b.second || (a.second == b.second && a.first > b.first);
        };
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(comp)> maxHeap(comp);

        for (const auto& it : num2cnt) {
            int num = it.first, cnt = it.second;
            maxHeap.push({num, cnt});
            if (maxHeap.size() > k) {
                maxHeap.pop();
            }
        }

        vector<int> res(k);
        int i = k - 1;
        while (maxHeap.size()) {
            res[i] = maxHeap.top().first;
            maxHeap.pop();
            i --;
        }

        return res;
    }
};

/*
    快速排序的选择方法

    时间复杂度：O(nlogn)
        n 为数组的长度。
        处理的过程包括一次遍历和一次子分支的递归，平均时间复杂度为O(nlogn)。
        最坏情况下，每次取的基准元素都位于数组的两端，时间复杂度退化为 O(n^2)。
    空间复杂度：O(n)
        哈希表的大小为 O(n)，数组的大小为 O(k)，共计为 O(n)。
*/
class Solution_3 {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> num2cnt;
        for (auto& v: nums) {
            num2cnt[v] ++;
        }

        vector<pair<int, int>> values;
        for (auto& kv : num2cnt) {
            values.push_back(kv);
        }

        vector<int> res;
        quickSort(values, 0, values.size() - 1, res, k);
        return res;
    }
     
    void quickSort(vector<pair<int, int>>& v, int start, int end, vector<int>& res, int k) {
        int picked = rand() % (end - start + 1) + start;
        swap(v[picked], v[start]);

        int pivot = v[start].second;
        int index = start;
        for (int i = start + 1; i <= end; i ++) {
            if (v[i].second >= pivot) {
                swap(v[index + 1], v[i]);
                index ++;
            }
        }
        swap(v[start], v[index]);

        if (k <= index - start) {
            quickSort(v, start, index - 1, res, k);
        } else {
            for (int i = start; i <= index; i++) {
                res.push_back(v[i].first);
            }
            if (k > index - start + 1) {
                quickSort(v, index + 1, end, res, k - (index - start + 1));
            }
        }
    }
};

int main() {
    Solution solution;
    vector<int> nums = {1, 1, 1, 2, 2, 3};
    printArray(nums);
    int k = 2;
    vector<int> topK = solution.topKFrequent(nums, k);
    cout << "The " << k << " highest frequency elements are: ";
    for (int num : topK) {
        cout << num << " ";
    }
    cout << endl;

    return 0;
}

Golang 代码

package main

import (
	"container/heap"
	"fmt"
)

/*
   统计频率：
       使用哈希表统计每个元素的出现次数。
   最小堆：
       使用最小堆（优先队列）维护前 K 个高频元素。
       当堆的大小超过 K 时，弹出堆顶元素（频率最小的元素）。
   提取结果：
       从堆中依次弹出元素，得到前 K 个高频元素。

   时间复杂度：O(nlogk)
       统计频率：O(n)，其中 n 是数组的长度。
       堆操作：每次插入和弹出操作的时间复杂度为 O(logk)，总共操作 n 次，因此总时间复杂度为 O(nlogk)。
   空间复杂度：O(n)
       哈希表：O(n)，用于存储每个元素的频率。
       堆：O(k)，用于存储前 K 个高频元素。
*/
// topKFrequent 返回数组中出现频率最高的 k 个元素
func topKFrequent(nums []int, k int) []int {
	// 统计每个数字的出现频率
	freq := make(map[int]int)
	for _, num := range nums {
		freq[num] ++
	}

	// 创建最小堆
	h := &MinHeap{}
	heap.Init(h)

	// 将元素推入堆并维护堆的大小
	for num, count := range freq {
		heap.Push(h, Pair{value: num, count: count})
		if h.Len() > k {
			heap.Pop(h)
		}
	}

	// 从堆中提取结果
	result := make([]int, k)
	for i := k - 1; i >= 0; i -- {
		result[i] = heap.Pop(h).(Pair).value
	}
	return result
}

// Pair 定义一个键值对结构
type Pair struct {
	value int // 元素值
	count int // 出现次数
}

// MinHeap 是一个最小堆的实现
type MinHeap []Pair

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

func (h *MinHeap) Push(x interface{}) {
	*h = append(*h, x.(Pair))
}

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

func main() {
	// 测试用例
	testCases := []struct {
		nums    []int
		k       int
		expected []int
	}{
		{[]int{1, 1, 1, 2, 2, 3}, 2, []int{1, 2}},
		{[]int{1}, 1, []int{1}},
	}

	for i, tc := range testCases {
		result := topKFrequent(tc.nums, tc.k)
		fmt.Printf("Test Case %d, Input: nums = %v, k = %d\n", i+1, tc.nums, tc.k)

		if fmt.Sprint(result) == fmt.Sprint(tc.expected) {
			fmt.Printf("Test Case %d, Output: %v, PASS\n", i+1, result)
		} else {
			fmt.Printf("Test Case %d, Output: %v, FAIL (Expected: %v)\n", i+1, result, tc.expected)
		}
	}
}