85 lines
2.3 KiB
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
|
||
|
}
|