优先队列在C++/Java等语言中已经存在于标准库中,而go语言中却没有。今天用了下,留个纪念吧
package main
import (
"container/heap"
"fmt"
)
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
}
func TestIntHeap(t *testing.T) {
pq := &IntHeap{}
heap.Init(pq)
for i := 1; i < 10; i++ {
heap.Push(pq, i)
}
for len(*pq) > 0 {
x := heap.Pop(pq)
fmt.Println(x)
}
}
应用:leetcode 5710 积压订单中的订单总数
这道题目耗费了挺长时间,主要原因是对heap构造优先列队的使用不熟练,Less函数的具体含义不清楚,Top函数应该返回第0个元素还是最后一个元素,也弄混淆了。
import (
"container/heap"
)
type buyHeap [][]int
type sellHeap [][]int
// buyHeap 为最大堆,每次Pop返回的最大值,取Top的话应该是数组中的第0号元素
func (h buyHeap) Len() int { return len(h) }
func (h buyHeap) Less(i, j int) bool { return h[i][0] > h[j][0] }
func (h buyHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *buyHeap) Push(x interface{}) {
*h = append(*h, x.([]int))
}
func (h buyHeap) Top() interface{} {
if h.Len() == 0 {
return nil
}
return h[0]
}
func (h *buyHeap) Pop() interface{} {
old := *h
n := len(old)
x := old[n-1]
*h = old[0 : n-1]
return x
}
// sellHeap 和 buyHeap 相反
func (h sellHeap) Len() int { return len(h) }
func (h sellHeap) Less(i, j int) bool { return h[i][0] < h[j][0] }
func (h sellHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *sellHeap) Push(x interface{}) {
*h = append(*h, x.([]int))
}
func (h sellHeap) Top() interface{} {
if h.Len() == 0 {
return nil
}
return h[0]
}
func (h *sellHeap) Pop() interface{} {
old := *h
n := len(old)
x := old[n-1]
*h = old[0 : n-1]
return x
}
func getNumberOfBacklogOrders(orders [][]int) int {
const M = 1000000000 + 7
buys := &buyHeap{}
sells := &sellHeap{}
heap.Init(buys)
heap.Init(sells)
for i := range orders {
o := orders[i]
pr, am, typ := o[0], o[1], o[2]
if typ == 0 {
// buy
if sells.Len() == 0 || sells.Top().([]int)[0] > pr {
heap.Push(buys, o)
} else {
for sells.Len() > 0 && am > 0 {
if sells.Top().([]int)[0] > pr {
break
}
tp := heap.Pop(sells).([]int)
if am >= tp[1] {
am -= tp[1]
} else {
tp[1] -= am
am = 0
heap.Push(sells, tp)
}
}
if am > 0 {
heap.Push(buys, []int{pr, am, typ})
}
}
} else {
// sell
if buys.Len() == 0 || buys.Top().([]int)[0] < pr {
heap.Push(sells, o)
} else {
for buys.Len() > 0 && am > 0 {
if buys.Top().([]int)[0] < pr {
break
}
tp := heap.Pop(buys).([]int)
if am >= tp[1] {
am -= tp[1]
} else {
tp[1] -= am
am = 0
heap.Push(buys, tp)
}
}
if am > 0 {
heap.Push(sells, []int{pr, am, typ})
}
}
}
}
ans := 0
// 显示类型转换
buysInts := [][]int(*buys)
sellInts := [][]int(*sells)
for _, o := range buysInts {
ans = (ans + o[1]) % M
}
for _, o := range sellInts {
ans = (ans + o[1]) % M
}
return ans
}