114 lines
1.9 KiB
Go
114 lines
1.9 KiB
Go
package heap
|
|
|
|
import (
|
|
"cmp"
|
|
"container/heap"
|
|
"slices"
|
|
)
|
|
|
|
type Heap[T comparable] struct {
|
|
Data []T
|
|
Index map[T]map[int]bool
|
|
lessFunc func(a, b T) bool
|
|
}
|
|
|
|
func More[T cmp.Ordered](x, y T) bool {
|
|
return !cmp.Less(x, y) && x != y
|
|
}
|
|
|
|
func NewBy[T comparable](lessFunc func(a, b T) bool, items ...T) *Heap[T] {
|
|
h := &Heap[T]{slices.Clone(items), map[T]map[int]bool{}, lessFunc}
|
|
heap.Init(h)
|
|
|
|
for i, v := range h.Data {
|
|
if h.Index[v] == nil {
|
|
h.Index[v] = map[int]bool{}
|
|
}
|
|
|
|
h.Index[v][i] = true
|
|
}
|
|
|
|
return h
|
|
}
|
|
|
|
func New[T cmp.Ordered](items ...T) *Heap[T] {
|
|
return NewBy(func(a, b T) bool { return a < b }, items...)
|
|
}
|
|
|
|
func (h *Heap[T]) putIndex(item T, index int) {
|
|
if h.Index == nil {
|
|
h.Index = map[T]map[int]bool{}
|
|
}
|
|
|
|
if h.Index[item] == nil {
|
|
h.Index[item] = map[int]bool{}
|
|
}
|
|
|
|
h.Index[item][index] = true
|
|
}
|
|
|
|
func (h *Heap[T]) dropIndex(item T, index int) {
|
|
if h.Index == nil {
|
|
h.Index = map[T]map[int]bool{}
|
|
}
|
|
|
|
if h.Index[item] == nil {
|
|
h.Index[item] = map[int]bool{}
|
|
}
|
|
|
|
delete(h.Index[item], index)
|
|
}
|
|
|
|
func (h Heap[T]) Len() int { return len(h.Data) }
|
|
func (h Heap[T]) Less(i, j int) bool { return h.lessFunc(h.Data[i], h.Data[j]) }
|
|
func (h Heap[T]) Swap(i, j int) {
|
|
h.dropIndex(h.Data[i], i)
|
|
h.dropIndex(h.Data[j], j)
|
|
|
|
h.Data[i], h.Data[j] = h.Data[j], h.Data[i]
|
|
|
|
h.putIndex(h.Data[i], i)
|
|
h.putIndex(h.Data[j], j)
|
|
}
|
|
|
|
func (h *Heap[T]) Push(x any) {
|
|
datum := x.(T)
|
|
h.Data = append(h.Data, datum)
|
|
h.putIndex(datum, len(h.Data)-1)
|
|
}
|
|
|
|
func (h *Heap[T]) Pop() any {
|
|
index := len(h.Data) - 1
|
|
datum := h.Data[index]
|
|
h.dropIndex(datum, index)
|
|
|
|
h.Data = h.Data[:index]
|
|
return datum
|
|
}
|
|
|
|
func (h *Heap[T]) GetIndex(item T) int {
|
|
if h.Index[item] == nil {
|
|
return -1
|
|
}
|
|
|
|
var key int
|
|
for k := range h.Index[item] {
|
|
key = k
|
|
break
|
|
}
|
|
|
|
return key
|
|
}
|
|
|
|
func (h *Heap[T]) Drop(item T) {
|
|
heap.Remove(h, h.GetIndex(item))
|
|
}
|
|
|
|
func (h *Heap[T]) Put(item T) {
|
|
heap.Push(h, item)
|
|
}
|
|
|
|
func (h *Heap[T]) Top() T {
|
|
return h.Data[0]
|
|
}
|