github.com/djordje200179/extendedlibrary/datastructures@v1.7.1-0.20240227175559-d09520a92dd4/seqs/pq/queue.go (about)

     1  package pq
     2  
     3  import (
     4  	"github.com/djordje200179/extendedlibrary/datastructures/iter"
     5  	"github.com/djordje200179/extendedlibrary/misc/functions/comparison"
     6  )
     7  
     8  // Queue is a priority queue implementation based on a binary heap.
     9  // By default, the queue is a min-heap, but a custom comparator can be provided.
    10  type Queue[T any] struct {
    11  	slice []T
    12  
    13  	comparator comparison.Comparator[T]
    14  }
    15  
    16  // New creates a new priority queue with the given comparator.
    17  func New[T any](comparator comparison.Comparator[T]) *Queue[T] {
    18  	pq := &Queue[T]{
    19  		slice: make([]T, 0),
    20  
    21  		comparator: comparator,
    22  	}
    23  
    24  	return pq
    25  }
    26  
    27  // NewFromIterable creates a new Queue with the given comparator and elements from the given iter.Iterable
    28  func NewFromIterable[T any](comparator comparison.Comparator[T], iterable iter.Iterable[T]) *Queue[T] {
    29  	pq := New[T](comparator)
    30  
    31  	for it := iterable.Iterator(); it.Valid(); it.Move() {
    32  		pq.PushBack(it.Get())
    33  	}
    34  
    35  	return pq
    36  }
    37  
    38  // Empty returns true if the queue is empty.
    39  func (pq *Queue[T]) Empty() bool {
    40  	return len(pq.slice) == 0
    41  }
    42  
    43  // PushBack pushes a new element into the queue.
    44  // The element is inserted at the end of the queue and then moved up the heap until the heap property is satisfied.
    45  func (pq *Queue[T]) PushBack(value T) {
    46  	pq.slice = append(pq.slice, value)
    47  
    48  	currNode := len(pq.slice) - 1
    49  	for {
    50  		parentNode := (currNode - 1) / 2
    51  		if parentNode == currNode || pq.comparator(pq.slice[currNode], pq.slice[parentNode]) != comparison.FirstSmaller {
    52  			break
    53  		}
    54  
    55  		pq.slice[parentNode], pq.slice[currNode] = pq.slice[currNode], pq.slice[parentNode]
    56  		currNode = parentNode
    57  	}
    58  }
    59  
    60  // PeekFront returns the element at the front of the queue.
    61  // Panics if the queue is empty.
    62  func (pq *Queue[T]) PeekFront() T {
    63  	if pq.Empty() {
    64  		panic("Priority queue is empty")
    65  	}
    66  
    67  	return pq.slice[0]
    68  }
    69  
    70  // PopFront removes and returns the element at the front of the queue.
    71  // Panics if the queue is empty.
    72  func (pq *Queue[T]) PopFront() T {
    73  	if pq.Empty() {
    74  		panic("Priority queue is empty")
    75  	}
    76  
    77  	lastIndex := len(pq.slice) - 1
    78  
    79  	pq.slice[0], pq.slice[lastIndex] = pq.slice[lastIndex], pq.slice[0]
    80  
    81  	currNode := 0
    82  	for {
    83  		leftChild := 2*currNode + 1
    84  		if leftChild >= lastIndex || leftChild < 0 {
    85  			break
    86  		}
    87  
    88  		child := leftChild
    89  		if rightChild := leftChild + 1; rightChild < lastIndex && pq.comparator(pq.slice[rightChild], pq.slice[leftChild]) == comparison.FirstSmaller {
    90  			child = rightChild
    91  		}
    92  
    93  		if pq.comparator(pq.slice[child], pq.slice[currNode]) != comparison.FirstSmaller {
    94  			break
    95  		}
    96  
    97  		pq.slice[currNode], pq.slice[child] = pq.slice[child], pq.slice[currNode]
    98  		currNode = child
    99  	}
   100  
   101  	lastElem := pq.slice[lastIndex]
   102  	pq.slice = pq.slice[:lastIndex]
   103  	return lastElem
   104  }