feat: progress

pkg/tools/heap/main.go

@@ -3,10 +3,12 @@ package heap
 import (
 	"cmp"
 	"container/heap"
+	"slices"
 )
 
-type Heap[T any] struct {
+type Heap[T comparable] struct {
 	Data     []T
+	Index    map[T]map[int]bool
 	lessFunc func(a, b T) bool
 }
 
@@ -14,22 +16,98 @@ func More[T cmp.Ordered](x, y T) bool {
 	return !cmp.Less(x, y) && x != y
 }
 
-func NewBy[T any](lessFunc func(a, b T) bool, items ...T) *Heap[T] {
-	h := &Heap[T]{items, lessFunc}
+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[T](func(a, b T) bool { return a < b }, items...)
+	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.Data[i], h.Data[j] = h.Data[j], h.Data[i] }
+func (h Heap[T]) Swap(i, j int) {
+	h.dropIndex(h.Data[i], i)
+	h.dropIndex(h.Data[j], j)
 
-func (h *Heap[T]) Push(x any) { h.Data = append(h.Data, x.(T)) }
-func (h *Heap[T]) Pop() any {
-	defer func() { h.Data = h.Data[:len(h.Data)-1] }()
-	return h.Data[len(h.Data)-1]
+	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]
 }