94 lines
1.8 KiB
Go
94 lines
1.8 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
|
||
|
|
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) {
|
||
|
|
// TODO: increase by square when capacity is full
|
||
|
|
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)
|
||
|
|
pq.heap = pq.heap[:pq.n]
|
||
|
|
pq.n--
|
||
|
|
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(start int) {
|
||
|
|
for k := start; k > 1 && pq.less(pq.heap[k/2], pq.heap[k]); k = k / 2 {
|
||
|
|
pq.swap(k/2, k)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (pq *priorityQueue[T]) sink(k int) {
|
||
|
|
// While k is a parent with some children
|
||
|
|
for 2*k <= pq.n {
|
||
|
|
|
||
|
|
// First child of a parent k
|
||
|
|
j := 2 * k
|
||
|
|
|
||
|
|
// If first child is less than second
|
||
|
|
// we choose second one for exchange.
|
||
|
|
// Parent should be swapped with biggest child
|
||
|
|
if j < pq.n && pq.less(pq.heap[j], pq.heap[j+1]) {
|
||
|
|
j++
|
||
|
|
}
|
||
|
|
|
||
|
|
// If parent is already bigger than biggest child
|
||
|
|
// we found the position
|
||
|
|
if !pq.less(pq.heap[k], pq.heap[j]) {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
|
||
|
|
// swap parent node with child
|
||
|
|
pq.swap(j, k)
|
||
|
|
|
||
|
|
// child node is a new parent
|
||
|
|
k = j
|
||
|
|
}
|
||
|
|
}
|