1
+ module heap
2
+ import resizable_array
3
+
4
+ /// Resizable 2-ary min-heap, backed by a resizable array
5
+ /// `cmp` defines the ordering of elements
6
+ record Heap[T](rawContents: ResizableArray[T], cmp: (T, T) => Ordering at {})
7
+
8
+ /// Make a new Heap with the given comparison operation
9
+ def heap[T](cmp: (T,T) => Ordering at {}) =
10
+ Heap[T](resizableArray(), cmp)
11
+
12
+ /// Make a new Heap with the given comparison operation and initial capacity
13
+ def heap[T](cmp: (T,T) => Ordering at {}, capacity: Int) =
14
+ Heap[T](resizableArray(capacity), cmp)
15
+
16
+ namespace internal {
17
+ def left(idx: Int) = 2 * idx + 1
18
+ def right(idx: Int) = 2 * idx + 2
19
+ def parent(idx: Int) = (idx - 1) / 2
20
+
21
+ def bubbleUp[A](heap: Heap[A], idx: Int) = {
22
+ val arr = heap.rawContents
23
+ arr.boundsCheck(idx) // idx > parent(idx), parent(parent(idx)) etc
24
+
25
+ def go(idx: Int): Unit = {
26
+ if (idx > 0 and (heap.cmp)(arr.unsafeGet(parent(idx)), arr.unsafeGet(idx)) is Greater()) {
27
+ arr.unsafeSwap(parent(idx), idx)
28
+ go(parent(idx))
29
+ }
30
+ }
31
+ go(idx)
32
+ }
33
+ def sinkDown[A](heap: Heap[A], idx: Int) = {
34
+ val arr = heap.rawContents
35
+
36
+ def infixLt(x: A, y: A) = {
37
+ if ((heap.cmp)(x,y) is Less()) { true } else { false }
38
+ }
39
+
40
+ def go(idx: Int): Unit = {
41
+ if (right(idx) < arr.size) {
42
+ val v = arr.unsafeGet(idx)
43
+ val l = arr.unsafeGet(left(idx))
44
+ val r = arr.unsafeGet(right(idx))
45
+ if (l < v && r < v) {
46
+ // swap with the smaller one
47
+ if (l < r) {
48
+ arr.unsafeSwap(left(idx), idx)
49
+ go(left(idx))
50
+ } else {
51
+ arr.unsafeSwap(right(idx), idx)
52
+ go(right(idx))
53
+ }
54
+ } else if (l < v) {
55
+ arr.unsafeSwap(left(idx), idx)
56
+ go(left(idx))
57
+ } else if (r < v) {
58
+ arr.unsafeSwap(right(idx), idx)
59
+ go(right(idx))
60
+ }
61
+ } else if (left(idx) < arr.size) {
62
+ if (arr.unsafeGet(left(idx)) < arr.unsafeGet(idx)) {
63
+ arr.unsafeSwap(left(idx), idx)
64
+ go(left(idx))
65
+ }
66
+ } // else: we are at the bottom
67
+ }
68
+ go(idx)
69
+ }
70
+ }
71
+
72
+ /// Insert value into heap
73
+ ///
74
+ /// O(log n) worst case if capacity suffices, O(1) average
75
+ def insert[T](heap: Heap[T], value: T): Unit = {
76
+ with on[OutOfBounds].panic();
77
+ val idx = heap.rawContents.add(value)
78
+ internal::bubbleUp(heap, idx)
79
+ }
80
+
81
+ /// find and return (but not remove) the minimal element in this heap
82
+ ///
83
+ /// O(1)
84
+ def findMin[T](heap: Heap[T]): T / Exception[OutOfBounds] = {
85
+ heap.rawContents.get(0)
86
+ }
87
+
88
+ /// find and remove the minimal element in this heap
89
+ ///
90
+ /// O(log n)
91
+ def deleteMin[T](heap: Heap[T]): T / Exception[OutOfBounds] = {
92
+ val res = heap.rawContents.get(0)
93
+ heap.rawContents.unsafeSet(0, heap.rawContents.popRight())
94
+ internal::sinkDown(heap, 0)
95
+ res
96
+ }
97
+
98
+ /// Number of elements in the heap
99
+ ///
100
+ /// O(1)
101
+ def size[T](heap: Heap[T]): Int = {
102
+ heap.rawContents.size
103
+ }
0 commit comments