A heap in Go

heap.go

package heap

type Heap struct {
	data    []int
	compare func(int, int) bool
}

func New() *Heap {
	return &Heap{
		data:    []int{-1},
		compare: func(a, b int) bool { return a > b },
	}
}

func (heap *Heap) Insert(value int) {
	heap.data = append(heap.data, value)
	heap.swim(len(heap.data) - 1)
}

func (heap *Heap) Top() int {
	return heap.data[1]
}

func (heap *Heap) Pop() (top int) {
	top = heap.data[1]
	heap.swap(1, len(heap.data)-1)
	heap.data = heap.data[:len(heap.data)-1]
	heap.sink(1)

	return top
}

func (heap *Heap) Size() int {
	return len(heap.data) - 1
}

func (heap *Heap) IsEmpty() bool {
	return len(heap.data) == 1
}

func (heap *Heap) shouldBeHigher(first, second int) bool {
	return heap.compare(heap.data[first], heap.data[second])
}

func parentIndex(index int) int {
	if index == 1 {
		return index
	} else {
		return index / 2
	}
}

func leftChildIndex(index int) int {
	return 2 * index
}

func rightChildIndex(index int) int {
	return 2*index + 1
}

func (heap *Heap) swap(first, second int) {
	heap.data[first], heap.data[second] = heap.data[second], heap.data[first]
}

func (heap *Heap) swim(index int) {
	for index != 1 {
		parent := parentIndex(index)
		if heap.shouldBeHigher(index, parent) {
			heap.swap(parent, index)
			index = parent
		} else {
			break
		}
	}
}

func (heap *Heap) sink(index int) {
	for index < len(heap.data) {
		leftChild := leftChildIndex(index)
		if leftChild >= len(heap.data) {
			break
		}

		betterChild := leftChild
		rightChild := leftChild + 1
		rightOk := rightChild < len(heap.data)
		if rightOk && heap.shouldBeHigher(rightChild, leftChild) {
			betterChild = rightChild
		}

		if heap.shouldBeHigher(betterChild, index) {
			heap.swap(betterChild, index)
			index = betterChild
		} else {
			break
		}
	}
}