algorithms/sorting/priority_queue.go

package sorting

type PriorityQueue[T any] interface {
	top() T
	insert(T)
	delete() T
	size() int
	isEmpty() bool
}

type priorityQueue[T any] struct {
	n    int
	heap []T
	// Actually this is "less" function.
	// First element not necessary is less than second,
	// Can be used both ways to get reversed priority
	less func(T, T) bool
}

func NewPQ[T any](less func(T, T) bool) PriorityQueue[T] {
	return &priorityQueue[T]{
		n:    0,
		less: less,
		// First element is not used
		// to make equation "next = current * 2" work
		heap: make([]T, 1, 1),
	}
}

func (pq *priorityQueue[T]) top() T {
	// After all operations "biggest" element
	// should be always at the beginning of the array
	return pq.heap[1]
}

func (pq *priorityQueue[T]) insert(item T) {
	// We can ignore "resize optimization":
	// append function will create new array for new slice
	// with doubled capacity when needed
	// https://github.com/golang/go/blob/master/src/runtime/slice.go#L188
	// https://go.dev/blog/slices-intro
	// https://go.dev/play/p/OKtCFskbp2t
	// https://stackoverflow.com/questions/23531737/how-the-slice-is-enlarged-by-append-is-the-capacity-always-doubled
	pq.heap = append(pq.heap, item)
	pq.n++
	pq.swim(pq.n)
}

func (pq *priorityQueue[T]) delete() T {
	top := pq.top()
	pq.swap(1, pq.n)
	// Discard "biggest" element from queue
	// for possible GC
	pq.heap = pq.heap[:pq.n]
	pq.n--
	// Reshape heap from top to bottom
	pq.sink(1)
	return top
}

func (pq *priorityQueue[_]) size() int {
	return pq.n
}

func (pq *priorityQueue[_]) isEmpty() bool {
	return pq.n == 0
}

func (pq *priorityQueue[_]) swap(i, j int) {
	pq.heap[i], pq.heap[j] = pq.heap[j], pq.heap[i]
}

func (pq *priorityQueue[T]) swim(child int) {
	// Until we reach top of the heap
	// and parent node is less than current child
	for child > 1 && pq.less(pq.heap[child/2], pq.heap[child]) {

		// We swap parent with the child
		pq.swap(child/2, child)

		// Parent node becomes new child
		// for next iteration
		child = child / 2
	}
}

func (pq *priorityQueue[T]) sink(parent int) {
	// While parent has some children
	for 2*parent <= pq.n {

		// First child of a parent
		child := 2 * parent

		// If first child is less than second
		// we choose second one for exchange.
		// Parent should be swapped with biggest child
		if child < pq.n && pq.less(pq.heap[child], pq.heap[child+1]) {
			child++
		}

		// If parent is already bigger than biggest child
		// we found the position
		if !pq.less(pq.heap[parent], pq.heap[child]) {
			break
		}

		// swap parent node with child
		pq.swap(parent, child)

		// child node is a new parent
		parent = child
	}
}