v / vlib / datatypes
Raw file | 84 loc (76 sloc) | 2.38 KB | Latest commit hash 7d8c38672
1module datatypes
2
3// MinHeap is a binary minimum heap data structure.
4pub struct MinHeap[T] {
5mut:
6 data []T
7}
8
9// insert adds an element to the heap.
10pub fn (mut heap MinHeap[T]) insert(item T) {
11 // push item to the end of the array
12 heap.data << item
13 // swap the new node with its parent until the heap is in order
14 mut child := heap.data.len - 1
15 mut parent := heap.parent(child)
16 for heap.data[parent] > heap.data[child] {
17 heap.data[parent], heap.data[child] = heap.data[child], heap.data[parent]
18 child = parent
19 parent = heap.parent(child)
20 }
21}
22
23// pop removes the top-most element from the heap.
24pub fn (mut heap MinHeap[T]) pop() !T {
25 if heap.data.len == 0 {
26 return error('Heap is empty')
27 } else if heap.data.len == 1 {
28 return heap.data.pop()
29 }
30 item := heap.data[0]
31 // move last element to root
32 heap.data[0] = heap.data.pop()
33 // swap the new root with its minimum child until the heap is in order
34 mut parent := 0
35 mut left := heap.left_child(parent) or { return item }
36 mut right := heap.right_child(parent) or { left }
37 for heap.data[parent] > heap.data[left] || heap.data[parent] > heap.data[right] {
38 // choose min for min heap
39 swap := if heap.data[left] <= heap.data[right] { left } else { right }
40 heap.data[parent], heap.data[swap] = heap.data[swap], heap.data[parent]
41 parent = swap
42 left = heap.left_child(parent) or { break }
43 right = heap.right_child(parent) or { left }
44 }
45 return item
46}
47
48// peek gets the top-most element from the heap without removing it.
49pub fn (heap MinHeap[T]) peek() !T {
50 if heap.data.len == 0 {
51 return error('Heap is empty')
52 }
53 return heap.data[0]
54}
55
56// len returns the number of elements in the heap.
57pub fn (heap MinHeap[T]) len() int {
58 return heap.data.len
59}
60
61// left_child is a helper function that returns the index of the left
62// child given a parent idx, or none if there is no left child.
63fn (heap MinHeap[T]) left_child(idx int) !int {
64 child := 2 * idx + 1
65 if child >= heap.data.len {
66 return error('none')
67 }
68 return child
69}
70
71// right_child is a helper function that returns the index of the right
72// child given a parent idx, or none if there is no right child.
73fn (heap MinHeap[T]) right_child(idx int) !int {
74 child := 2 * idx + 2
75 if child >= heap.data.len {
76 return error('none')
77 }
78 return child
79}
80
81// parent is a helper function that returns the parent index of the child.
82fn (heap MinHeap[T]) parent(idx int) int {
83 return (idx - 1) / 2
84}