184 lines
4.6 KiB
Go
184 lines
4.6 KiB
Go
package top100liked
|
||
|
||
import (
|
||
"reflect"
|
||
"sort"
|
||
"testing"
|
||
)
|
||
|
||
// https://leetcode.cn/problems/top-k-frequent-elements/?envType=study-plan-v2&envId=top-100-liked
|
||
// 347. 前 K 个高频元素
|
||
//
|
||
// 给一个整数数组 nums 和整数 k，返回出现频率前 k 高的元素，顺序任意。
|
||
// 进阶要求时间复杂度优于 O(n log n)。
|
||
//
|
||
// 通用第一步：统计频率 map[int]int。下面三种解法都基于此。
|
||
|
||
// countFreq 统计每个元素出现次数。
|
||
func countFreq(nums []int) map[int]int {
|
||
m := make(map[int]int)
|
||
for _, v := range nums {
|
||
m[v]++
|
||
}
|
||
return m
|
||
}
|
||
|
||
// 解法一：最小堆 — O(n log k) / O(n)
|
||
// 维护大小为 k 的最小堆，堆顶是当前堆内频率最小者。
|
||
// 逐个把 (元素, 频率) 入堆；堆大小超过 k 时弹出堆顶，淘汰频率最小的。
|
||
// 遍历结束，堆里剩下的 k 个即前 k 高频元素。
|
||
//
|
||
// 实现要点（container/heap）：
|
||
// - 定义 type entry struct{ val, cnt int}
|
||
// - 定义 type minHeap []entry，实现 heap.Interface
|
||
// - Less 按 cnt 升序（堆顶最小）
|
||
// - 遍历 m：heap.Push(h, e)；len(*h) > k 时 heap.Pop(h)
|
||
// - 收集堆中剩余元素返回
|
||
type entry struct {
|
||
val int
|
||
cnt int
|
||
}
|
||
|
||
type minHeap struct {
|
||
data []entry
|
||
}
|
||
|
||
// 上浮：新元素放末尾后，与父比较，比父小则交换向上
|
||
// i 是需要 siftUp 的元素的 index
|
||
func (h *minHeap) siftUp(i int) {
|
||
for i > 0 {
|
||
p := (i - 1) / 2
|
||
if h.data[p].cnt > h.data[i].cnt {
|
||
h.data[p], h.data[i] = h.data[i], h.data[p]
|
||
i = p
|
||
} else {
|
||
break
|
||
}
|
||
}
|
||
}
|
||
|
||
// 下沉：从 i 开始，与较小子比较，比子大则交换向下。n 是有效堆大小
|
||
func (h *minHeap) siftDown(i, n int) {
|
||
for {
|
||
l, r := 2*i+1, 2*i+2
|
||
smallest := i
|
||
if l < n && h.data[l].cnt < h.data[smallest].cnt {
|
||
smallest = l
|
||
}
|
||
if r < n && h.data[r].cnt < h.data[smallest].cnt {
|
||
smallest = r
|
||
}
|
||
if smallest == i {
|
||
break
|
||
}
|
||
h.data[smallest], h.data[i] = h.data[i], h.data[smallest]
|
||
i = smallest
|
||
}
|
||
}
|
||
|
||
// push：追加末尾 + 上浮
|
||
func (h *minHeap) push(e entry) {
|
||
idx := len(h.data)
|
||
h.data = append(h.data, e)
|
||
h.siftUp(idx)
|
||
}
|
||
|
||
// popMin 弹出堆顶（最小 cnt）：用末尾覆盖堆顶，缩容，下沉
|
||
func (h *minHeap) popMin() entry {
|
||
if len(h.data) == 0 {
|
||
return entry{}
|
||
}
|
||
ret := h.data[0]
|
||
h.data[0] = h.data[len(h.data)-1]
|
||
h.data = h.data[:len(h.data)-1]
|
||
if len(h.data) > 0 {
|
||
h.siftDown(0, len(h.data))
|
||
}
|
||
return ret
|
||
}
|
||
|
||
func (h *minHeap) len() int { return len(h.data) }
|
||
|
||
func topKFrequentHeap(nums []int, k int) []int {
|
||
m := countFreq(nums)
|
||
h := minHeap{}
|
||
for v, c := range m {
|
||
h.push(entry{v, c})
|
||
if h.len() > k {
|
||
h.popMin()
|
||
}
|
||
}
|
||
res := make([]int, 0, k)
|
||
for h.len() > 0 {
|
||
res = append(res, h.popMin().val)
|
||
}
|
||
return res
|
||
}
|
||
|
||
// 解法二：桶排序 — O(n) 平均 / O(n)
|
||
// 频率取值范围 [0, n]，开长度 n+1 的桶：buckets[f] 存所有频率为 f 的元素。
|
||
// 从 buckets[n] 倒序向 buckets[1] 遍历，依次把桶内元素加入结果，凑够 k 个停止。
|
||
//
|
||
// 实现要点：
|
||
// - buckets := make([][]int, len(nums)+1)
|
||
// - for val, cnt := range m { buckets[cnt] = append(buckets[cnt], val) }
|
||
// - for f := len(buckets)-1; f >= 0 && len(res) < k; f-- { res = append(res, buckets[f]...) }
|
||
func topKFrequentBucket(nums []int, k int) []int {
|
||
m := countFreq(nums)
|
||
buckets := make([][]int, len(nums)+1)
|
||
for v, c := range m {
|
||
buckets[c] = append(buckets[c], v)
|
||
}
|
||
res := make([]int, 0, k)
|
||
for f := len(buckets) - 1; f >= 0 && len(res) < k; f-- {
|
||
res = append(res, buckets[f]...)
|
||
}
|
||
return res
|
||
}
|
||
|
||
// equalIgnoreOrder 忽略顺序比较两个 []int。
|
||
func equalIgnoreOrder(a, b []int) bool {
|
||
if len(a) != len(b) {
|
||
return false
|
||
}
|
||
ca := append([]int(nil), a...)
|
||
cb := append([]int(nil), b...)
|
||
sort.Ints(ca)
|
||
sort.Ints(cb)
|
||
return reflect.DeepEqual(ca, cb)
|
||
}
|
||
|
||
func TestTopKFrequent(t *testing.T) {
|
||
cases := []struct {
|
||
name string
|
||
nums []int
|
||
k int
|
||
want []int
|
||
}{
|
||
{"basic", []int{1, 1, 1, 2, 2, 3}, 2, []int{1, 2}},
|
||
{"single", []int{1}, 1, []int{1}},
|
||
{"tie", []int{1, 2}, 2, []int{1, 2}}, // 频率相同，都可返回
|
||
{"negative", []int{4, 1, -1, 2, -1, 2, 3}, 2, []int{-1, 2}},
|
||
{"duplicate candidates", []int{3, 0, 1, 0}, 1, []int{0}},
|
||
}
|
||
solvers := []struct {
|
||
name string
|
||
fn func([]int, int) []int
|
||
}{
|
||
{"heap", topKFrequentHeap},
|
||
{"bucket", topKFrequentBucket},
|
||
}
|
||
for _, s := range solvers {
|
||
t.Run(s.name, func(t *testing.T) {
|
||
for _, c := range cases {
|
||
// 复制 nums 防止解法修改原切片
|
||
nums := append([]int(nil), c.nums...)
|
||
got := s.fn(nums, c.k)
|
||
if !equalIgnoreOrder(got, c.want) {
|
||
t.Errorf("%s/%s: got %v, want %v", s.name, c.name, got, c.want)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|