advent-of-code-go/data-structures/heap.go

package structures

// MinHeap is an implementation of a min heap
type MinHeap struct {
	heap
}

// NewMinHeap initializes a heap with a closerToRootFunction that simply
// returns true if the first arg is smaller than the second
func NewMinHeap() MinHeap {
	nestedHeap := heap{
		closerToRoot: func(val1, val2 int) bool {
			return val1 < val2
		},
	}
	return MinHeap{nestedHeap}
}

// MaxHeap is an implementation of max heap
type MaxHeap struct {
	heap
}

// NewMaxHeap initializes a heap with a closerToRootFunction that simply
// returns true if the first arg is larger than the second
func NewMaxHeap() MaxHeap {
	nestedHeap := heap{
		closerToRoot: func(val1, val2 int) bool {
			return val1 > val2
		},
	}
	return MaxHeap{nestedHeap}
}

// heap contains a slice of heapNodes
// A heap can be represented as an array/slice with no gaps because
// calculating the indices of two children or the parent is simple
// from any given index
type heap struct {
	nodes        []heapNode
	closerToRoot func(val1, val2 int) bool
}

// heapNode is an interface making the type for a Min/MaxHeap node flexible
// nodes must be be able to state their value to be sorted by
type heapNode interface {
	Value() int
}

// Add appends a new node onto the heap and heapifies it
// to ensure correct ordering
func (h *heap) Add(newNode heapNode) {
	h.nodes = append(h.nodes, newNode)
	h.heapifyFromEnd()
}

// Remove returns the node at the root, i.e. the minimum value node
func (h *heap) Remove() heapNode {
	if len(h.nodes) == 0 {
		return nil
	}

	rootNode := h.nodes[0]

	// move last node to start & reduce length by one
	h.nodes[0] = h.nodes[len(h.nodes)-1]
	h.nodes = h.nodes[:len(h.nodes)-1]

	// heapify the heap from the start to sort the minimum value into the 0 index
	h.heapifyFromStart()

	return rootNode
}

func (h *heap) swap(i, j int) {
	h.nodes[i], h.nodes[j] = h.nodes[j], h.nodes[i]
}

// heapify from end expects an unordered value in the last index, it will compare
// it to its parent index and swapped if applicable, and repeated until the heap
// is valid
func (h *heap) heapifyFromEnd() {
	currentIndex := len(h.nodes) - 1
	for currentIndex > 0 {
		parentIndex := (currentIndex - 1) / 2
		parentNode := h.nodes[parentIndex]
		if h.closerToRoot(h.nodes[currentIndex].Value(), parentNode.Value()) {
			h.swap(parentIndex, currentIndex)
			currentIndex = parentIndex
		} else {
			break
		}
	}
}

// heapify from start expects an unordered value in the heap in index zero,
// that node's value is compared to its children, and swaps are made as needed
// until the heap is valid
func (h *heap) heapifyFromStart() {
	currentIndex := 0

	for {
		// find smaller of two children
		smallerChildIndex := currentIndex
		for i := 1; i <= 2; i++ {
			childIndex := currentIndex*2 + i
			// if a child value is closer to the root than the current node,
			// store it's index
			if childIndex < len(h.nodes) &&
				h.closerToRoot(h.nodes[childIndex].Value(), h.nodes[smallerChildIndex].Value()) {
				smallerChildIndex = childIndex
			}
		}

		// if smallerChildIndex was not reassigned, no swap is needed, return out
		if smallerChildIndex == currentIndex {
			return
		}

		// otherwise swap & update currentIndex to keep checking on next loop
		h.swap(smallerChildIndex, currentIndex)
		currentIndex = smallerChildIndex
	}
}