diff --git a/datastructure/heap/maxheap.go b/datastructure/heap/maxheap.go index 445f27a..d1e1d58 100644 --- a/datastructure/heap/maxheap.go +++ b/datastructure/heap/maxheap.go @@ -23,7 +23,13 @@ func NewMaxHeap[T any](comparator lancetconstraints.Comparator) *MaxHeap[T] { } } -// heapifyUp heapify the data from bottom to up +// Push value into the heap +func (h *MaxHeap[T]) Push(value T) { + h.data = append(h.data, value) + h.heapifyUp(len(h.data) - 1) +} + +// heapifyUp heapify the data from bottom to top func (h *MaxHeap[T]) heapifyUp(i int) { for h.comparator.Compare(h.data[parentIndex(i)], h.data[i]) < 0 { h.swap(parentIndex(i), i) @@ -31,6 +37,70 @@ func (h *MaxHeap[T]) heapifyUp(i int) { } } +// Pop return the largest value, and remove it from the heap +// if heap is empty, return zero value and fasle +func (h *MaxHeap[T]) Pop(value T) (T, bool) { + var val T + if h.Size() == 0 { + return val, false + } + + val = h.data[0] + l := len(h.data) - 1 + + h.data[0] = h.data[l] + h.data = h.data[:l] + h.heapifyDown(0) + + return val, true +} + +// heapifyDown heapify the data from top to bottom +func (h *MaxHeap[T]) heapifyDown(i int) { + lastIndex := len(h.data) - 1 + l, r := leftChildIndex(i), rightChildIndex(i) + childToCompare := 0 + + for l <= lastIndex { + if l == lastIndex { + childToCompare = l + } else if h.comparator.Compare(h.data[l], h.data[r]) > 0 { + childToCompare = l + } else { + childToCompare = r + } + + if h.comparator.Compare(h.data[i], h.data[childToCompare]) < 0 { + h.swap(i, childToCompare) + i = childToCompare + l, r = leftChildIndex(i), rightChildIndex(i) + } else { + break + } + } +} + +// Peek returns the largest element from the heap without removing it. +// if heap is empty, it returns zero value and false. +func (h *MaxHeap[T]) Peek() (T, bool) { + if h.Size() == 0 { + var val T + return val, false + } + + return h.data[0], true +} + +// Size return the number of elements in the heap +func (h *MaxHeap[T]) Size() int { + return len(h.data) +} + +// Data return data of the heap +func (h *MaxHeap[T]) Data() []T { + return h.data +} + // parentIndex get parent index of the given index func parentIndex(i int) int { return (i - 1) / 2