github.com/weedge/lib@v0.0.0-20230424045628-a36dcc1d90e4/container/queue/priority_queue.go (about)

     1  package queue
     2  
     3  import (
     4  	"container/heap"
     5  )
     6  
     7  type Item struct {
     8  	Value    interface{}
     9  	Priority int64
    10  	Index    int
    11  }
    12  
    13  // this is a priority queue as implemented by a min heap
    14  // ie. the 0th element is the *lowest* value
    15  type PriorityQueue []*Item
    16  
    17  func NewPriorityQueue(capacity int) PriorityQueue {
    18  	return make(PriorityQueue, 0, capacity)
    19  }
    20  
    21  func (pq PriorityQueue) Len() int {
    22  	return len(pq)
    23  }
    24  
    25  // min heap
    26  func (pq PriorityQueue) Less(i, j int) bool {
    27  	return pq[i].Priority < pq[j].Priority
    28  }
    29  
    30  func (pq PriorityQueue) Swap(i, j int) {
    31  	pq[i], pq[j] = pq[j], pq[i]
    32  	pq[i].Index = i
    33  	pq[j].Index = j
    34  }
    35  
    36  // heap push interface impl
    37  func (pq *PriorityQueue) Push(x interface{}) {
    38  	n := len(*pq)
    39  	c := cap(*pq)
    40  	if n+1 > c {
    41  		npq := make(PriorityQueue, n, c*2)
    42  		copy(npq, *pq)
    43  		*pq = npq
    44  	}
    45  	*pq = (*pq)[0 : n+1]
    46  	item := x.(*Item)
    47  	item.Index = n
    48  	(*pq)[n] = item
    49  }
    50  
    51  // heap pop interface impl
    52  func (pq *PriorityQueue) Pop() interface{} {
    53  	n := len(*pq)
    54  	c := cap(*pq)
    55  	if n < (c/2) && c > 25 {
    56  		npq := make(PriorityQueue, n, c/2)
    57  		copy(npq, *pq)
    58  		*pq = npq
    59  	}
    60  	item := (*pq)[n-1]
    61  	item.Index = -1
    62  	*pq = (*pq)[0 : n-1]
    63  	return item
    64  }
    65  
    66  func (pq *PriorityQueue) PeekAndShift(max int64) (*Item, int64) {
    67  	if pq.Len() == 0 {
    68  		return nil, 0
    69  	}
    70  
    71  	item := (*pq)[0]
    72  	if item.Priority > max {
    73  		return nil, item.Priority - max
    74  	}
    75  	heap.Remove(pq, 0)
    76  
    77  	return item, 0
    78  }
    79  
    80  func (pq *PriorityQueue) Top() (top *Item) {
    81  	if pq.Len() == 0 {
    82  		return
    83  	}
    84  
    85  	top = (*pq)[0]
    86  
    87  	return
    88  }
    89  
    90  // update item priority
    91  func (pq *PriorityQueue) Update(item *Item, value interface{}, priority int64) {
    92  	item.Value = value
    93  	item.Priority = priority
    94  	heap.Fix(pq, item.Index)
    95  }