144 lines
2.6 KiB
Go
144 lines
2.6 KiB
Go
package common
|
|
|
|
import (
|
|
"container/heap"
|
|
"sync"
|
|
)
|
|
|
|
/*
|
|
Example usage:
|
|
|
|
func main() {
|
|
h := NewHeap()
|
|
|
|
h.Push(String("msg1"), 1)
|
|
h.Push(String("msg3"), 3)
|
|
h.Push(String("msg2"), 2)
|
|
|
|
fmt.Println(h.Pop())
|
|
fmt.Println(h.Pop())
|
|
fmt.Println(h.Pop())
|
|
}
|
|
*/
|
|
|
|
type Comparable interface {
|
|
Less(o interface{}) bool
|
|
}
|
|
|
|
//-----------------------------------------------------------------------------
|
|
|
|
type Heap struct {
|
|
pq priorityQueue
|
|
}
|
|
|
|
func NewHeap() *Heap {
|
|
return &Heap{pq: make([]*pqItem, 0)}
|
|
}
|
|
|
|
func (h *Heap) Len() int64 {
|
|
return int64(len(h.pq))
|
|
}
|
|
|
|
func (h *Heap) Push(value interface{}, priority Comparable) {
|
|
heap.Push(&h.pq, &pqItem{value: value, priority: priority})
|
|
}
|
|
|
|
func (h *Heap) Peek() interface{} {
|
|
if len(h.pq) == 0 {
|
|
return nil
|
|
}
|
|
return h.pq[0].value
|
|
}
|
|
|
|
func (h *Heap) Pop() interface{} {
|
|
item := heap.Pop(&h.pq).(*pqItem)
|
|
return item.value
|
|
}
|
|
|
|
//-----------------------------------------------------------------------------
|
|
|
|
type CHeap struct {
|
|
mtx sync.Mutex
|
|
pq priorityQueue
|
|
}
|
|
|
|
func NewCHeap() *CHeap {
|
|
return &CHeap{pq: make([]*pqItem, 0)}
|
|
}
|
|
|
|
func (h *CHeap) Len() int64 {
|
|
h.mtx.Lock()
|
|
defer h.mtx.Unlock()
|
|
return int64(len(h.pq))
|
|
}
|
|
|
|
func (h *CHeap) Push(value interface{}, priority Comparable) {
|
|
h.mtx.Lock()
|
|
defer h.mtx.Unlock()
|
|
heap.Push(&h.pq, &pqItem{value: value, priority: priority})
|
|
}
|
|
|
|
func (h *CHeap) Peek() interface{} {
|
|
h.mtx.Lock()
|
|
defer h.mtx.Unlock()
|
|
if len(h.pq) == 0 {
|
|
return nil
|
|
}
|
|
return h.pq[0].value
|
|
}
|
|
|
|
func (h *CHeap) Pop() interface{} {
|
|
h.mtx.Lock()
|
|
defer h.mtx.Unlock()
|
|
item := heap.Pop(&h.pq).(*pqItem)
|
|
return item.value
|
|
}
|
|
|
|
//-----------------------------------------------------------------------------
|
|
|
|
///////////////////////
|
|
// From: http://golang.org/pkg/container/heap/#example__priorityQueue
|
|
|
|
type pqItem struct {
|
|
value interface{}
|
|
priority Comparable
|
|
index int
|
|
}
|
|
|
|
type priorityQueue []*pqItem
|
|
|
|
func (pq priorityQueue) Len() int { return len(pq) }
|
|
|
|
func (pq priorityQueue) Less(i, j int) bool {
|
|
return pq[i].priority.Less(pq[j].priority)
|
|
}
|
|
|
|
func (pq priorityQueue) Swap(i, j int) {
|
|
pq[i], pq[j] = pq[j], pq[i]
|
|
pq[i].index = i
|
|
pq[j].index = j
|
|
}
|
|
|
|
func (pq *priorityQueue) Push(x interface{}) {
|
|
n := len(*pq)
|
|
item := x.(*pqItem)
|
|
item.index = n
|
|
*pq = append(*pq, item)
|
|
}
|
|
|
|
func (pq *priorityQueue) Pop() interface{} {
|
|
old := *pq
|
|
n := len(old)
|
|
item := old[n-1]
|
|
item.index = -1 // for safety
|
|
*pq = old[0 : n-1]
|
|
return item
|
|
}
|
|
|
|
func (pq *priorityQueue) Update(item *pqItem, value interface{}, priority Comparable) {
|
|
heap.Remove(pq, item.index)
|
|
item.value = value
|
|
item.priority = priority
|
|
heap.Push(pq, item)
|
|
}
|