1 | module datatypes |
2 | |
3 | // MinHeap is a binary minimum heap data structure. |
4 | pub struct MinHeap[T] { |
5 | mut: |
6 | data []T |
7 | } |
8 | |
9 | // insert adds an element to the heap. |
10 | pub 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. |
24 | pub 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. |
49 | pub 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. |
57 | pub 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. |
63 | fn (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. |
73 | fn (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. |
82 | fn (heap MinHeap[T]) parent(idx int) int { |
83 | return (idx - 1) / 2 |
84 | } |