github.com/benz9527/toy-box/algo@v0.0.0-20240221120937-66c0c6bd5abd/queue/priority_queue.go (about)

     1  package queue
     2  
     3  import (
     4  	"container/heap"
     5  	"sync/atomic"
     6  )
     7  
     8  // 使用最小堆来构建优先级队列
     9  // 每次最小堆的根节点都会与树中最后一个叶子节点交换位置
    10  // 之后剪去该叶子节点就获得了当前优先级最高的信息
    11  // 剪枝的操作完成后，原本被提升为根节点的节点会被重新下沉
    12  // 完成优先级的重新分布
    13  // 新加入节点也会重新排列分布
    14  
    15  type pqItem[E comparable] struct {
    16  	priority   int64       // alignment size: 8; size: 8
    17  	index      int64       // alignment size: 8; size: 8
    18  	comparator LessThan[E] // alignment size: 8; size: 8
    19  	value      E           // based on the type of E
    20  }
    21  
    22  func (item *pqItem[E]) GetPriority() int64 {
    23  	return atomic.LoadInt64(&item.priority)
    24  }
    25  
    26  func (item *pqItem[E]) SetPriority(priority int64) {
    27  	atomic.StoreInt64(&item.priority, priority)
    28  }
    29  
    30  func (item *pqItem[E]) GetIndex() int64 {
    31  	return atomic.LoadInt64(&item.index)
    32  }
    33  
    34  func (item *pqItem[E]) SetIndex(index int64) {
    35  	atomic.SwapInt64(&item.index, index)
    36  }
    37  
    38  func (item *pqItem[E]) GetValue() E {
    39  	return item.value
    40  }
    41  
    42  func (item *pqItem[E]) GetLessThan() LessThan[E] {
    43  	return item.comparator
    44  }
    45  
    46  func (item *pqItem[E]) SetLessThan(comparator LessThan[E]) {
    47  	if comparator == nil {
    48  		// default comparator, default min priority value has high priority
    49  		comparator = func(i, j PQItem[E]) bool {
    50  			return i.GetPriority() < j.GetPriority()
    51  		}
    52  	}
    53  	item.comparator = comparator
    54  }
    55  
    56  func NewPQItem[T comparable](value T, priority int64) PQItem[T] {
    57  	return &pqItem[T]{
    58  		value:    value,
    59  		priority: priority,
    60  		index:    0,
    61  		comparator: func(i, j PQItem[T]) bool {
    62  			// default comparator, default min priority value has high priority
    63  			return i.GetPriority() < j.GetPriority()
    64  		},
    65  	}
    66  }
    67  
    68  type arrayPQ[E comparable] []PQItem[E]
    69  
    70  func (pq *arrayPQ[E]) Len() int {
    71  	return len(*pq)
    72  }
    73  
    74  func (pq *arrayPQ[E]) Less(i, j int) bool {
    75  	// Nil panic at your own risk
    76  	return (*pq)[i].GetLessThan()((*pq)[i], (*pq)[j])
    77  }
    78  
    79  func (pq *arrayPQ[E]) Swap(i, j int) {
    80  	(*pq)[i], (*pq)[j] = (*pq)[j], (*pq)[i]
    81  	(*pq)[i].SetIndex(int64(i))
    82  	(*pq)[j].SetIndex(int64(j))
    83  }
    84  
    85  func (pq *arrayPQ[E]) Pop() interface{} {
    86  	prev := *pq
    87  	n := len(prev)
    88  	if n <= 0 {
    89  		return nil
    90  	}
    91  
    92  	item := prev[n-1]
    93  	item.SetIndex(-1)
    94  	prev[n-1] = *new(PQItem[E]) // nil object
    95  	*pq = prev[:n-1]
    96  	return item
    97  }
    98  
    99  func (pq *arrayPQ[E]) Push(i interface{}) {
   100  	item, ok := i.(PQItem[E])
   101  	if !ok {
   102  		return
   103  	}
   104  
   105  	prev := *pq
   106  	n := len(prev)
   107  	item.SetIndex(int64(n))
   108  	*pq = append(*pq, item.(PQItem[E]))
   109  }
   110  
   111  type ArrayPriorityQueue[E comparable] struct {
   112  	queue           arrayPQ[E]
   113  	localComparator LessThan[E]
   114  }
   115  
   116  // NewArrayPriorityQueue create a priority queue, you can customize the less than comparator.
   117  // Default less than comparator implemented by min priority value has high priority.
   118  func NewArrayPriorityQueue[E comparable](
   119  	capacity int,
   120  	comparator ...LessThan[E],
   121  ) PriorityQueue[E] {
   122  	if len(comparator) <= 0 {
   123  		comparator = []LessThan[E]{
   124  			func(i, j PQItem[E]) bool {
   125  				return i.GetPriority() < j.GetPriority()
   126  			},
   127  		}
   128  	}
   129  	if capacity <= 0 {
   130  		capacity = 32
   131  	}
   132  	pq := &ArrayPriorityQueue[E]{
   133  		queue:           make(arrayPQ[E], 0, capacity),
   134  		localComparator: comparator[0],
   135  	}
   136  	heap.Init(&pq.queue)
   137  	return pq
   138  }
   139  
   140  func (pq *ArrayPriorityQueue[E]) Len() int64 {
   141  	return int64(len(pq.queue))
   142  }
   143  
   144  func (pq *ArrayPriorityQueue[E]) Pop() (nilItem PQItem[E]) {
   145  	if len(pq.queue) == 0 {
   146  		return nil
   147  	}
   148  	item := heap.Pop(&pq.queue)
   149  	return item.(PQItem[E])
   150  }
   151  
   152  func (pq *ArrayPriorityQueue[E]) Push(item PQItem[E]) {
   153  	item.SetLessThan(pq.localComparator)
   154  	heap.Push(&pq.queue, item)
   155  }
   156  
   157  func (pq *ArrayPriorityQueue[E]) Peek() PQItem[E] {
   158  	if len(pq.queue) == 0 {
   159  		return nil
   160  	}
   161  	return pq.queue[0]
   162  }