github.com/chenjiandongx/go-queue@v0.0.0-20191023082232-e2a36f382f84/priority_queue.go (about) 1 package collections 2 3 import ( 4 "container/heap" 5 "sync" 6 ) 7 8 type PriorityQueue struct { 9 nodes []*PqNode 10 mut *sync.RWMutex 11 } 12 13 // PriorityQueue Node 14 type PqNode struct { 15 Value string 16 Priority, index int 17 } 18 19 func NewPriorityQueue() *PriorityQueue { 20 pq := &PriorityQueue{mut: new(sync.RWMutex)} 21 heap.Init(pq) 22 return pq 23 } 24 25 func (pq *PriorityQueue) Put(v *PqNode) { 26 defer pq.mut.Unlock() 27 pq.mut.Lock() 28 heap.Push(pq, v) 29 } 30 31 func (pq *PriorityQueue) Get() (interface{}, bool) { 32 defer pq.mut.Unlock() 33 pq.mut.Lock() 34 if len(pq.nodes) > 0 { 35 item := heap.Pop(pq) 36 return item, true 37 } 38 return nil, false 39 } 40 41 func (pq PriorityQueue) Qsize() int { 42 defer pq.mut.RUnlock() 43 pq.mut.RLock() 44 return len(pq.nodes) 45 } 46 47 func (pq *PriorityQueue) IsEmpty() bool { 48 defer pq.mut.RUnlock() 49 pq.mut.RLock() 50 return !(len(pq.nodes) > 0) 51 } 52 53 // `Sort` interface Len() 54 func (pq PriorityQueue) Len() int { 55 return len(pq.nodes) 56 } 57 58 // `Sort` interface Less() 59 func (pq PriorityQueue) Less(i, j int) bool { 60 return pq.nodes[i].Priority > pq.nodes[j].Priority 61 } 62 63 // `Sort` interface Swap() 64 func (pq PriorityQueue) Swap(i, j int) { 65 pq.nodes[i], pq.nodes[j] = pq.nodes[j], pq.nodes[i] 66 pq.nodes[i].index, pq.nodes[j].index = i, j 67 } 68 69 // `Heap` interface Push() 70 func (pq *PriorityQueue) Push(v interface{}) { 71 item := v.(*PqNode) 72 item.index = len(pq.nodes) 73 pq.nodes = append(pq.nodes, item) 74 heap.Fix(pq, item.index) 75 } 76 77 // `Heap` interface Pop() 78 func (pq *PriorityQueue) Pop() interface{} { 79 old := *pq 80 n := len(old.nodes) 81 item := old.nodes[n-1] 82 item.index = -1 83 pq.nodes = old.nodes[0 : n-1] 84 return item 85 }