From e873ac03ea650e2a86180a7e1bf22652989a79fe Mon Sep 17 00:00:00 2001
From: jonastheis <4181434+jonastheis@users.noreply.github.com>
Date: Tue, 29 Oct 2024 16:00:18 +0800
Subject: [PATCH] add HeapMap component

---
 common/heapmap.go | 90 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 90 insertions(+)
 create mode 100644 common/heapmap.go

diff --git a/common/heapmap.go b/common/heapmap.go
new file mode 100644
index 000000000000..90f51e2db25d
--- /dev/null
+++ b/common/heapmap.go
@@ -0,0 +1,90 @@
+package common
+
+type HeapMap[K comparable, T Comparable[T]] struct {
+	h              *Heap[T]
+	m              *ShrinkingMap[K, *HeapElement[T]]
+	keyFromElement func(T) K
+}
+
+func NewHeapMap[K comparable, T Comparable[T]](keyFromElement func(T) K) *HeapMap[K, T] {
+	return &HeapMap[K, T]{
+		h:              NewHeap[T](),
+		m:              NewShrinkingMap[K, *HeapElement[T]](1000),
+		keyFromElement: keyFromElement,
+	}
+}
+
+func (hm *HeapMap[K, T]) Len() int {
+	return hm.h.Len()
+}
+
+func (hm *HeapMap[K, T]) Push(element T) bool {
+	k := hm.keyFromElement(element)
+
+	if hm.m.Has(k) {
+		return false
+	}
+
+	heapElement := hm.h.Push(element)
+	hm.m.Set(k, heapElement)
+
+	return true
+}
+
+func (hm *HeapMap[K, T]) Pop() T {
+	element := hm.h.Pop()
+	k := hm.keyFromElement(element.Value())
+	hm.m.Delete(k)
+
+	return element.Value()
+}
+
+func (hm *HeapMap[K, T]) Peek() T {
+	return hm.h.Peek().Value()
+}
+
+func (hm *HeapMap[K, T]) RemoveByElement(element T) bool {
+	key := hm.keyFromElement(element)
+	heapElement, exists := hm.m.Get(key)
+	if !exists {
+		return false
+	}
+
+	hm.h.Remove(heapElement)
+	hm.m.Delete(key)
+
+	return true
+}
+
+func (hm *HeapMap[K, T]) RemoveByKey(key K) bool {
+	heapElement, exists := hm.m.Get(key)
+	if !exists {
+		return false
+	}
+
+	hm.h.Remove(heapElement)
+	hm.m.Delete(key)
+
+	return true
+}
+
+func (hm *HeapMap[K, T]) Clear() {
+	hm.h.Clear()
+	hm.m = NewShrinkingMap[K, *HeapElement[T]](1000)
+}
+
+func (hm *HeapMap[K, T]) Keys() []K {
+	return hm.m.Keys()
+}
+
+func (hm *HeapMap[K, T]) Elements() []T {
+	var elements []T
+	for _, element := range hm.m.Values() {
+		elements = append(elements, element.Value())
+	}
+	return elements
+}
+
+func (hm *HeapMap[K, T]) Has(element T) bool {
+	return hm.m.Has(hm.keyFromElement(element))
+}