github.com/songzhibin97/go-baseutils@v0.0.2-0.20240302024150-487d8ce9c082/structure/trees/binaryheap/binaryheap.go

// Package binaryheap implements a binary heap backed by array list.
//
// Comparator defines this heap as either min or max heap.
//
// Structure is not thread safe.
//
// References: http://en.wikipedia.org/wiki/Binary_heap
package binaryheap

import (
	"fmt"
	"github.com/songzhibin97/go-baseutils/base/bcomparator"
	"github.com/songzhibin97/go-baseutils/structure/lists/arraylist"
	"github.com/songzhibin97/go-baseutils/structure/trees"

	"strings"
)

// Assert Tree implementation
var _ trees.Tree[any] = (*Heap[any])(nil)

// Heap holds elements in an array-list
type Heap[E any] struct {
	list       *arraylist.List[E]
	Comparator bcomparator.Comparator[E]
}

// NewWith instantiates a new empty heap tree with the custom comparator.
func NewWith[E any](comparator bcomparator.Comparator[E]) *Heap[E] {
	return &Heap[E]{list: arraylist.New[E](), Comparator: comparator}
}

// NewWithIntComparator instantiates a new empty heap with the IntComparator, i.e. elements are of type int.
func NewWithIntComparator() *Heap[int] {
	return &Heap[int]{list: arraylist.New[int](), Comparator: bcomparator.IntComparator()}
}

// NewWithStringComparator instantiates a new empty heap with the StringComparator, i.e. elements are of type string.
func NewWithStringComparator() *Heap[string] {
	return &Heap[string]{list: arraylist.New[string](), Comparator: bcomparator.StringComparator()}
}

// Push adds a value onto the heap and bubbles it up accordingly.
func (heap *Heap[E]) Push(values ...E) {
	if len(values) == 1 {
		heap.list.Add(values[0])
		heap.bubbleUp()
	} else {
		// Reference: https://en.wikipedia.org/wiki/Binary_heap#Building_a_heap
		for _, value := range values {
			heap.list.Add(value)
		}
		size := heap.list.Size()/2 + 1
		for i := size; i >= 0; i-- {
			heap.bubbleDownIndex(i)
		}
	}
}

// Pop removes top element on heap and returns it, or nil if heap is empty.
// Second return parameter is true, unless the heap was empty and there was nothing to pop.
func (heap *Heap[E]) Pop() (value E, ok bool) {
	value, ok = heap.list.Get(0)
	if !ok {
		return
	}
	lastIndex := heap.list.Size() - 1
	heap.list.Swap(0, lastIndex)
	heap.list.Remove(lastIndex)
	heap.bubbleDown()
	return
}

// Peek returns top element on the heap without removing it, or nil if heap is empty.
// Second return parameter is true, unless the heap was empty and there was nothing to peek.
func (heap *Heap[E]) Peek() (value E, ok bool) {
	return heap.list.Get(0)
}

// Empty returns true if heap does not contain any elements.
func (heap *Heap[E]) Empty() bool {
	return heap.list.Empty()
}

// Size returns number of elements within the heap.
func (heap *Heap[E]) Size() int {
	return heap.list.Size()
}

// Clear removes all elements from the heap.
func (heap *Heap[E]) Clear() {
	heap.list.Clear()
}

// Values returns all elements in the heap.
func (heap *Heap[E]) Values() []E {
	values := make([]E, heap.list.Size(), heap.list.Size())
	for it := heap.Iterator(); it.Next(); {
		values[it.Index()] = it.Value()
	}
	return values
}

// String returns a string representation of container
func (heap *Heap[E]) String() string {
	b := strings.Builder{}
	b.WriteString("BinaryHeap\n")
	for it := heap.Iterator(); it.Next(); {
		b.WriteString(fmt.Sprintf("(value:%v) ", it.Value()))
	}
	return b.String()
}

// Performs the "bubble down" operation. This is to place the element that is at the root
// of the heap in its correct place so that the heap maintains the min/max-heap order property.
func (heap *Heap[E]) bubbleDown() {
	heap.bubbleDownIndex(0)
}

// Performs the "bubble down" operation. This is to place the element that is at the index
// of the heap in its correct place so that the heap maintains the min/max-heap order property.
func (heap *Heap[E]) bubbleDownIndex(index int) {
	size := heap.list.Size()
	for leftIndex := index<<1 + 1; leftIndex < size; leftIndex = index<<1 + 1 {
		rightIndex := index<<1 + 2
		smallerIndex := leftIndex
		leftValue, _ := heap.list.Get(leftIndex)
		rightValue, _ := heap.list.Get(rightIndex)
		if rightIndex < size && heap.Comparator(leftValue, rightValue) > 0 {
			smallerIndex = rightIndex
		}
		indexValue, _ := heap.list.Get(index)
		smallerValue, _ := heap.list.Get(smallerIndex)
		if heap.Comparator(indexValue, smallerValue) > 0 {
			heap.list.Swap(index, smallerIndex)
		} else {
			break
		}
		index = smallerIndex
	}
}

// Performs the "bubble up" operation. This is to place a newly inserted
// element (i.e. last element in the list) in its correct place so that
// the heap maintains the min/max-heap order property.
func (heap *Heap[E]) bubbleUp() {
	index := heap.list.Size() - 1
	for parentIndex := (index - 1) >> 1; index > 0; parentIndex = (index - 1) >> 1 {
		indexValue, _ := heap.list.Get(index)
		parentValue, _ := heap.list.Get(parentIndex)
		if heap.Comparator(parentValue, indexValue) <= 0 {
			break
		}
		heap.list.Swap(index, parentIndex)
		index = parentIndex
	}
}

// Check that the index is within bounds of the list
func (heap *Heap[E]) withinRange(index int) bool {
	return index >= 0 && index < heap.list.Size()
}

// UnmarshalJSON @implements json.Unmarshaler
func (heap *Heap[E]) UnmarshalJSON(bytes []byte) error {
	return heap.list.UnmarshalJSON(bytes)
}

// MarshalJSON @implements json.Marshaler
func (heap *Heap[E]) MarshalJSON() ([]byte, error) {
	return heap.list.MarshalJSON()
}