113 lines
2.6 KiB
Go
113 lines
2.6 KiB
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
|
|
}
|
|
}
|