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  }