Skip to content

Commit

Permalink
add HeapMap component
Browse files Browse the repository at this point in the history
  • Loading branch information
jonastheis committed Oct 29, 2024
1 parent 113241e commit e873ac0
Showing 1 changed file with 90 additions and 0 deletions.
90 changes: 90 additions & 0 deletions common/heapmap.go
Original file line number Diff line number Diff line change
@@ -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))
}

0 comments on commit e873ac0

Please sign in to comment.