在Leetcode刷题的时候,有些题目需要自己实现一个堆的结构,在此记录一下。例如Leetcode-23:合并k个升序链表。
1. 首先实现一个链表节点的小顶堆的结构,按照值排序大小
// 实现小顶堆, Len, Less, Swap, Push, Pop 五方法
type NodeHeap []*ListNode
func (nh NodeHeap) Len() int { return len(nh) }
func (nh NodeHeap) Less(i,j int) bool { return nh[i].Val < nh[j].Val }
func (nh NodeHeap) Swap(i,j int) { nh[i], nh[j] = nh[j], nh[i] }
func (nh *NodeHeap) Push(x interface{}) {
*nh = append(*nh, x.(*ListNode))
}
func (nh *NodeHeap) Pop() interface{} {
old := *nh
// 弹出堆顶元素
n := len(old)
x := old[n - 1]
*nh = old[:n - 1]
return x
}
2. 初始化节点堆结构
nh := &NodeHeap{}
dummy := &ListNode{}
tail := dummy
heap.Init(nh)
for _, node := range lists {
if node != nil {
heap.Push(nh, node)
}
}
// 取出堆顶元素并且加到链表后面
for nh.Len() != 0 {
// 弹出, 用接口的方法
x := heap.Pop(nh).(*ListNode)
tail.Next = x
tail = tail.Next
// 加入
if x.Next != nil {
heap.Push(nh, x.Next)
}
}
return dummy.Next
通过初始化NodeHeap结构,然后调用heap.Init(nh)来将这个数组堆化。然后将各个节点调用heap.Push(nh, node)推入堆中,并自行堆化。然后从小顶堆中弹出最小值也就是堆顶元素的时候使用heap.Pop(nh),这里要使用类型断言,因为heap.Pop弹出的元素类型是空接口类型。
在一次周赛中遇到了一种新的堆的写法,在此补充一下:Leetcode-2462题,雇佣K位工人。这道题使用了两个小顶堆维护,然后通过双指针碰撞选出相对较小的元素,以下是代码解释:
func totalCost(costs []int, k int, candidates int) int64 {
// 两个小顶堆
n := len(costs)
res := 0
if candidates * 2 < n {
pre := hp{ costs[:candidates] }
heap.Init(&pre)
suf := hp{ costs[n - candidates:] }
heap.Init(&suf)
// 双指针碰撞
for i,j := candidates, n - candidates - 1; k > 0 && i <= j; k-- {
if pre.IntSlice[0] <= suf.IntSlice[0] {
res += pre.IntSlice[0]
pre.IntSlice[0] = costs[i]
heap.Fix(&pre, 0)
i++
} else {
res += suf.IntSlice[0]
suf.IntSlice[0] = costs[j]
heap.Fix(&suf, 0)
j--
}
}
// 不满足条件了 合并
costs = append(pre.IntSlice, suf.IntSlice...)
}
// 此时不满足candidates * 2 < n, 直接数组排序,取前k个
sort.Ints(costs)
for i := 0; i < k; i++ {
res += costs[i]
}
return int64(res)
}
type hp struct { sort.IntSlice }
func (h hp) Push(interface{}) {}
func (h hp) Pop() (_ interface{}) { return }
原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。
如有侵权,请联系 cloudcommunity@tencent.com 删除。
原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。
如有侵权,请联系 cloudcommunity@tencent.com 删除。