当前位置：首页 > news >正文

TopK问题（堆排序）-- go

news 2025/8/24 9:11:35

TopK问题（堆排序）

一、堆

1.1 堆的定义

完全二叉树 + 堆序性（指的是大顶堆/小顶堆）

大顶堆：节点本身的值为根的子树/树的各个节点的最大值（递归性理解：只需要节点大于两个子节点）

小顶堆：节点本身的值为根的子树/树的各个节点的最小值（递归性理解：只需要节点小于两个子节点）

1.2堆的存储

层序遍历后，可以放在一维数组/切片中，有切片索引和堆位置索引的规律：

节点索引为 i，其子节点为 2i+1、2i+2

1.3 堆的基操

上滤：节点向上移动（原因是下面的节点破环堆序性）-- 插入场景

下滤：节点向下移动（原因是上面的节点破环堆序性）-- 调整

1.4 建堆

自顶向下：保持顶部堆序性的建堆方法+ 破环底部 + 上滤

自底向上：保持底部堆序性的建堆方法+ 破环顶部 + 下滤（全插入再调整）

1.5 应用

堆排序：本质还是使用优先队列

优先队列：在go中实现heap接口（PList 为优先队列类型）之后可以使用heap包的相关方法

// 这里采用 int 类型作为元素
// topK问题：采用堆排序的方式（优先队列的数据结构）// 实现优先队列
type PList []int// 实现go内置的heap接口
// 修改切片内容，只需要使用值接收者
func (p PList) Len() int{return len(p)
}func (p PList) Swap(i,j int){p[i], p[j] = p[j], p[i]
}// 大顶堆
func (p PList) Less(i,j int) bool{return p[i] > p[j]
}
// 修改切片本身，需要使用指针接收者
func (p *PList) Push(x interface{}){// 解引用出旧的切片old := *p// 类型转换后加入*p = append(old, x.(int))
}func (p *PList) Pop() interface{}{// 解引用出旧的切片old := *pn := len(old)// 弹出最后一个元素x := old[n-1]*p = old[0:n-1]return x
}

二、TopK问题

面试题 17.14. 最小K个数 - 力扣（LeetCode）

整体思路：创建大顶堆，将大元素全部放到大顶堆中，heap堆

func smallestK(arr []int, k int) []int {// 边界：if k <= 0 || k > len(arr) {return nil}// 创建大顶堆h := make(PList, k)// 先放入k个元素copy(h, arr[:k])// 初始化为堆结构heap.Init(&h)// 遍历剩余元素for _,num := range arr[k:]{// 因为是大顶堆，新元素小则需要修复if num < h[0]{h[0] = num// 索引0处的修改，需要修复heap.Fix(&h, 0)}}// 将堆中元素取出并返回result := make([]int, k)for i := 0; i < k; i++ {result[i] = heap.Pop(&h).(int)}return result
}// topK问题：采用堆排序的方式（优先队列的数据结构）// 实现优先队列
type PList []int// 实现go内置的heap接口
// 修改切片内容，只需要使用值接收者
func (p PList) Len() int{return len(p)
}func (p PList) Swap(i,j int){p[i], p[j] = p[j], p[i]
}// 大顶堆
func (p PList) Less(i,j int) bool{return p[i] > p[j]
}
// 修改切片本身，需要使用指针接收者
func (p *PList) Push(x interface{}){// 解引用出旧的切片old := *p// 类型转换后加入*p = append(old, x.(int))
}
// 大顶堆
func (p *PList) Pop() interface{}{// 解引用出旧的切片old := *pn := len(old)// 弹出最后一个元素x := old[n-1]*p = old[0:n-1]return x
}

// 采用小顶堆
func smallestK(arr []int, k int) []int {if k <= 0 || k > len(arr) {return nil}h := make(PList, k)copy(h, arr[:k])heap.Init(&h)for _,num := range arr[k:]{// 小顶堆：只允许大元素加入if num > h[0]{h[0] = numheap.Fix(&h, 0)}}result := make([]int, k)for i := 0; i < k; i++ {result[i] = heap.Pop(&h).(int)}return result
}type PList []intfunc (p PList) Len() int{return len(p)
}func (p PList) Swap(i,j int){p[i], p[j] = p[j], p[i]
}// 小顶堆: 小于时返回true
func (p PList) Less(i,j int) bool{return p[i] < p[j]
}func (p *PList) Push(x interface{}){old := *p*p = append(old, x.(int))
}func (p *PList) Pop() interface{}{old := *pn := len(old)x := old[n-1]*p = old[0:n-1]return x
}