v/vlib/datatypes/heap.v

85 lines
2.3 KiB
V

module datatypes
// 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
}