实现代码
package core
import "container/heap"
type Item struct {
Value interface{}
Index int
Priority int
}
type PriorityQueue []*Item
func NewPriorityQueue(cap int) PriorityQueue {
return make(PriorityQueue, cap)
}
func (pg PriorityQueue) Len() int {
return len(pg)
}
func (pg PriorityQueue) Less(i, j int) bool {
return pg[i].Priority < pg[j].Priority
}
func (pg PriorityQueue) Swap(i, j int) {
pg[i], pg[j] = pg[j], pg[i]
pg[i].Index = i
pg[j].Index = j
}
func (pg *PriorityQueue) Push(x interface{}) {
n := len(*pg)
c := cap(*pg)
if n+1 > c {
npg := make(PriorityQueue, c*2)
copy(npg, *pg)
*pg = npg
}
*pg = (*pg)[0:n+1]
item := x.(*Item)
item.Index = n
(*pg)[n] = item
}
func (pg *PriorityQueue) Pop() interface{} {
n := len(*pg)
c := cap(*pg)
if n < (c/2) && c > 25 {
npg := make(PriorityQueue, n, c/2)
copy(npg, *pg)
*pg = npg
}
item := (*pg)[n-1]
item.Index = -1
*pg = (*pg)[0:n-1]
return item
}
func (pg *PriorityQueue) PeekAndShift(max int) (*Item, int) {
if pg.Len() == 0 {
return nil, 0
}
item := (*pg)[0]
if item.Priority > max {
return nil, item.Priority - max
}
heap.Remove(pg, 0)
return item, 0
}
测试代码
import (
"testing"
"container/heap"
"sort"
"math/rand"
"reflect"
"path/filepath"
"runtime"
)
func equal(t *testing.T, act, exp interface{}) {
if !reflect.DeepEqual(exp, act) {
_, file, line, _ := runtime.Caller(1)
t.Logf("\033[31m%s:%d:\n\n\texp: %#v\n\n\tgot: %#v\033[39m\n\n",
filepath.Base(file), line, exp, act)
t.FailNow()
}
}
func TestPriorityQueue(t *testing.T) {
c := 100
pq := NewPriorityQueue(c)
for i := 0; i < c+1; i++ {
heap.Push(&pq, &Item{Value: i, Priority: i})
}
equal(t, pq.Len(), c+1)
equal(t, cap(pq), c*2)
for i := 0; i < c+1; i++ {
item := heap.Pop(&pq)
equal(t, item.(*Item).Value.(int), i)
}
equal(t, cap(pq), c/4)
}
func TestUnsortedInsert(t *testing.T) {
c := 100
pq := NewPriorityQueue(c)
ints := make([]int, 0, c)
for i := 0; i < c; i++ {
v := rand.Int()
ints = append(ints, v)
heap.Push(&pq, &Item{Value: i, Priority: v})
}
equal(t, pq.Len(), c)
equal(t, cap(pq), c)
sort.Sort(sort.IntSlice(ints))
for i := 0; i < c; i++ {
item, _ := pq.PeekAndShift(ints[len(ints)-1])
equal(t, item.Priority, int64(ints[i]))
}
}
func TestRemove(t *testing.T) {
c := 100
pq := NewPriorityQueue(c)
for i := 0; i < c; i++ {
v := rand.Int()
heap.Push(&pq, &Item{Value: "test", Priority: v})
}
for i := 0; i < 10; i++ {
heap.Remove(&pq, rand.Intn((c-1)-i))
}
lastPriority := heap.Pop(&pq).(*Item).Priority
for i := 0; i < (c - 10 - 1); i++ {
item := heap.Pop(&pq)
equal(t, lastPriority < item.(*Item).Priority, true)
lastPriority = item.(*Item).Priority
}
}