122 lines
3.5 KiB
Go
122 lines
3.5 KiB
Go
// Copyright 2009 The Go Authors. All rights reserved.
|
|
// Use of this source code is governed by a BSD-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
// Package heap provides heap operations for any type that implements
|
|
// heap.Interface. A heap is a tree with the property that each node is the
|
|
// minimum-valued node in its subtree.
|
|
//
|
|
// The minimum element in the tree is the root, at index 0.
|
|
//
|
|
// A heap is a common way to implement a priority queue. To build a priority
|
|
// queue, implement the Heap interface with the (negative) priority as the
|
|
// ordering for the Less method, so Push adds items while Pop removes the
|
|
// highest-priority item from the queue. The Examples include such an
|
|
// implementation; the file example_pq_test.go has the complete source.
|
|
//
|
|
// This package is a copy of the Go standard library's
|
|
// container/heap, but using generics.
|
|
package heap
|
|
|
|
import "sort"
|
|
|
|
// The Interface type describes the requirements
|
|
// for a type using the routines in this package.
|
|
// Any type that implements it may be used as a
|
|
// min-heap with the following invariants (established after
|
|
// Init has been called or if the data is empty or sorted):
|
|
//
|
|
// !h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len()
|
|
//
|
|
// Note that Push and Pop in this interface are for package heap's
|
|
// implementation to call. To add and remove things from the heap,
|
|
// use heap.Push and heap.Pop.
|
|
type Interface[V any] interface {
|
|
sort.Interface
|
|
Push(x V) // add x as element Len()
|
|
Pop() V // remove and return element Len() - 1.
|
|
}
|
|
|
|
// Init establishes the heap invariants required by the other routines in this package.
|
|
// Init is idempotent with respect to the heap invariants
|
|
// and may be called whenever the heap invariants may have been invalidated.
|
|
// The complexity is O(n) where n = h.Len().
|
|
func Init[V any](h Interface[V]) {
|
|
// heapify
|
|
n := h.Len()
|
|
for i := n/2 - 1; i >= 0; i-- {
|
|
down(h, i, n)
|
|
}
|
|
}
|
|
|
|
// Push pushes the element x onto the heap.
|
|
// The complexity is O(log n) where n = h.Len().
|
|
func Push[V any](h Interface[V], x V) {
|
|
h.Push(x)
|
|
up(h, h.Len()-1)
|
|
}
|
|
|
|
// Pop removes and returns the minimum element (according to Less) from the heap.
|
|
// The complexity is O(log n) where n = h.Len().
|
|
// Pop is equivalent to Remove(h, 0).
|
|
func Pop[V any](h Interface[V]) V {
|
|
n := h.Len() - 1
|
|
h.Swap(0, n)
|
|
down(h, 0, n)
|
|
return h.Pop()
|
|
}
|
|
|
|
// Remove removes and returns the element at index i from the heap.
|
|
// The complexity is O(log n) where n = h.Len().
|
|
func Remove[V any](h Interface[V], i int) V {
|
|
n := h.Len() - 1
|
|
if n != i {
|
|
h.Swap(i, n)
|
|
if !down(h, i, n) {
|
|
up(h, i)
|
|
}
|
|
}
|
|
return h.Pop()
|
|
}
|
|
|
|
// Fix re-establishes the heap ordering after the element at index i has changed its value.
|
|
// Changing the value of the element at index i and then calling Fix is equivalent to,
|
|
// but less expensive than, calling Remove(h, i) followed by a Push of the new value.
|
|
// The complexity is O(log n) where n = h.Len().
|
|
func Fix[V any](h Interface[V], i int) {
|
|
if !down(h, i, h.Len()) {
|
|
up(h, i)
|
|
}
|
|
}
|
|
|
|
func up[V any](h Interface[V], j int) {
|
|
for {
|
|
i := (j - 1) / 2 // parent
|
|
if i == j || !h.Less(j, i) {
|
|
break
|
|
}
|
|
h.Swap(i, j)
|
|
j = i
|
|
}
|
|
}
|
|
|
|
func down[V any](h Interface[V], i0, n int) bool {
|
|
i := i0
|
|
for {
|
|
j1 := 2*i + 1
|
|
if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
|
|
break
|
|
}
|
|
j := j1 // left child
|
|
if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
|
|
j = j2 // = 2*i + 2 // right child
|
|
}
|
|
if !h.Less(j, i) {
|
|
break
|
|
}
|
|
h.Swap(i, j)
|
|
i = j
|
|
}
|
|
return i > i0
|
|
}
|