v/vlib/adt/heap.v

module adt

// MinHeap is a binary minimum heap data structure.
pub struct MinHeap<T> {
mut:
	data []T
}

// insert adds an element to the heap.
pub fn (mut heap MinHeap<T>) insert(item T) {
	// push item to the end of the array
	heap.data << item
	// swap the new node with its parent until the heap is in order
	mut child := heap.data.len - 1
	mut parent := heap.parent(child)
	for heap.data[parent] > heap.data[child] {
		heap.data[parent], heap.data[child] = heap.data[child], heap.data[parent]
		child = parent
		parent = heap.parent(child)
	}
}

// pop removes the top-most element from the heap.
pub fn (mut heap MinHeap<T>) pop() ?T {
	if heap.data.len == 0 {
		return none
	} else if heap.data.len == 1 {
		return heap.data.pop()
	}
	item := heap.data[0]
	// move last element to root
	heap.data[0] = heap.data.pop()
	// swap the new root with its minimum child until the heap is in order
	mut parent := 0
	mut left := heap.left_child(parent) or { return item }
	mut right := heap.right_child(parent) or { left }
	for heap.data[parent] > heap.data[left] || heap.data[parent] > heap.data[right] {
		// choose min for min heap
		swap := if heap.data[left] <= heap.data[right] { left } else { right }
		heap.data[parent], heap.data[swap] = heap.data[swap], heap.data[parent]
		parent = swap
		left = heap.left_child(parent) or { break }
		right = heap.right_child(parent) or { left }
	}
	return item
}

// peek gets the top-most element from the heap without removing it.
pub fn (heap MinHeap<T>) peek() ?T {
	if heap.data.len == 0 {
		return none
	}
	return heap.data[0]
}

// len returns the number of elements in the heap.
pub fn (heap MinHeap<T>) len() int {
	return heap.data.len
}

// left_child is a helper function that returns the index of the left
// child given a parent idx, or none if there is no left child.
fn (heap MinHeap<T>) left_child(idx int) ?int {
	child := 2 * idx + 1
	if child >= heap.data.len {
		return none
	}
	return child
}

// right_child is a helper function that returns the index of the right
// child given a parent idx, or none if there is no right child.
fn (heap MinHeap<T>) right_child(idx int) ?int {
	child := 2 * idx + 2
	if child >= heap.data.len {
		return none
	}
	return child
}

// parent is a helper function that returns the parent index of the child.
fn (heap MinHeap<T>) parent(idx int) int {
	return (idx - 1) / 2
}