go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/collections/priority_queue.go (about)

     1  /*
     2  
     3  Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.
     5  
     6  */
     7  
     8  package collections
     9  
    10  // NewPriorityQueue returns a new priority queue.
    11  //
    12  // A priority queue is very similar to a heap but
    13  // instead of bring intrinsically ordered, you can supply a
    14  // separate priority.
    15  func NewPriorityQueue[T any]() *PriorityQueue[T] {
    16  	return &PriorityQueue[T]{
    17  		heap: &Heap[PriorityQueueItem[T]]{
    18  			LessFn: func(i, j PriorityQueueItem[T]) bool {
    19  				return i.Priority < j.Priority
    20  			},
    21  		},
    22  	}
    23  }
    24  
    25  // PriorityQueue is a heap of items with priorities.
    26  type PriorityQueue[T any] struct {
    27  	heap *Heap[PriorityQueueItem[T]]
    28  }
    29  
    30  // PriorityQueueItem is an item in the priority queue.
    31  type PriorityQueueItem[T any] struct {
    32  	Item     T
    33  	Priority int
    34  }
    35  
    36  // Len returns the length of the priority queue.
    37  func (pq *PriorityQueue[T]) Len() int {
    38  	return pq.heap.Len()
    39  }
    40  
    41  // Push pushes an item into the priority queue.
    42  func (pq *PriorityQueue[T]) Push(item T, priority int) {
    43  	pq.heap.Push(PriorityQueueItem[T]{
    44  		Item:     item,
    45  		Priority: priority,
    46  	})
    47  }
    48  
    49  // Peek returns the minimum item and its priority but does not remove it.
    50  func (pq *PriorityQueue[T]) Peek() (item T, priority int, ok bool) {
    51  	var pi PriorityQueueItem[T]
    52  	pi, ok = pq.heap.Peek()
    53  	if !ok {
    54  		return
    55  	}
    56  	item = pi.Item
    57  	priority = pi.Priority
    58  	return
    59  }
    60  
    61  // Pop pops an item off the priority queue.
    62  func (pq *PriorityQueue[T]) Pop() (item T, priority int, ok bool) {
    63  	var pi PriorityQueueItem[T]
    64  	pi, ok = pq.heap.Pop()
    65  	if !ok {
    66  		return
    67  	}
    68  	item = pi.Item
    69  	priority = pi.Priority
    70  	return
    71  }
    72  
    73  // Push pushes an item into the priority queue.
    74  func (pq *PriorityQueue[T]) PushPop(inputItem T, inputPriority int) (outputItem T, outputPriority int, ok bool) {
    75  	var pi PriorityQueueItem[T]
    76  	pi, ok = pq.heap.PushPop(PriorityQueueItem[T]{
    77  		Item:     inputItem,
    78  		Priority: inputPriority,
    79  	})
    80  	if !ok {
    81  		return
    82  	}
    83  	outputItem = pi.Item
    84  	outputPriority = pi.Priority
    85  	return
    86  }