go利用最小堆实现优先队列

实现代码

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)
}
// 实现接口heap.Interface接口
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
}

// add x as element Len()
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
}

// remove and return element Len() - 1.
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
    }
}