优先队列在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
}