2024 day 18 Complete
This commit is contained in:
151
helpers/heap.go
Normal file
151
helpers/heap.go
Normal file
@@ -0,0 +1,151 @@
|
||||
package aoc
|
||||
|
||||
/* Heap
|
||||
* This is a heap implementation with slice backing.
|
||||
* The less function is used to compare elements.
|
||||
* The heap is a complete binary tree, which means:
|
||||
* - All levels are filled except possibly for the last level.
|
||||
* - The last level is filled as much as possible from left to right.
|
||||
*
|
||||
* This binary tree also should satisfy the heap property (we use min-heap here):
|
||||
* - The smallest element is at the root
|
||||
* - The value of each node is less than or equal to the values of its children
|
||||
*
|
||||
* To implement a heap, we use a slice to store the elements with these rules:
|
||||
* - Root is at 0
|
||||
* - Left child of node at index i is at 2*i+1
|
||||
* - Right child of node at index i is at 2*i+2
|
||||
* - Parent of node at index i is at (i-1)/2
|
||||
* For example, if we insert: 2, 0, 7, 4, 3, 6, 9
|
||||
* The heap will look like this:
|
||||
* i: 0 1 2 3 4 5 6
|
||||
* [ 0, 2, 6, 4, 3, 7, 9 ]
|
||||
* And the effective binary tree looks like this:
|
||||
* 0
|
||||
* / \
|
||||
* 2 6
|
||||
* / \ / \
|
||||
* 4 3 7 9
|
||||
* So the root is at index [0],
|
||||
* the children of [0] art at 2*0+1 = [1] and 2*0+2 = [2]
|
||||
* The parent of [1] is at (1-1)/2 = [0],
|
||||
* and the parent of [2] is at (2-1)/2 = [0]
|
||||
* The children of [1] are at 2*1+1 = [3] and 2*1+2 = [4]
|
||||
* The parent of [3] is at (3-1)/2 = [1],
|
||||
* and the parent of [4] is at (4-1)/2 = [1]
|
||||
* The children of [2] are at 2*2+1 = [5] and 2*2+2 = [6]
|
||||
* The parent of [5] is at (5-1)/2 = [2],
|
||||
* and the parent of [6] is at (6-1)/2 = [2]
|
||||
|
||||
* Thanks to insomnes (https://github.com/insomnes) for this.
|
||||
*/
|
||||
|
||||
type Heap[T any] struct {
|
||||
data []T
|
||||
less func(a, b T) bool
|
||||
}
|
||||
|
||||
func NewHeap[T any](less func(a, b T) bool) *Heap[T] {
|
||||
return &Heap[T]{data: []T{}, less: less}
|
||||
}
|
||||
|
||||
func (h *Heap[T]) Len() int {
|
||||
return len(h.data)
|
||||
}
|
||||
|
||||
// Push adds a new element to the heap (the end of the slice)
|
||||
// Then restores the heap property
|
||||
func (h *Heap[T]) Push(x T) {
|
||||
h.data = append(h.data, x)
|
||||
h.ascend(len(h.data) - 1)
|
||||
}
|
||||
|
||||
// Pop removes the smallest element from the heap (first in the slice)
|
||||
// Under the hood, it swaps the first and last elements, pops the last element,
|
||||
// and then restores the heap property by comparing the new psuedo-root with
|
||||
// children
|
||||
func (h *Heap[T]) Pop() T {
|
||||
if len(h.data) == 0 {
|
||||
panic("heap is empty")
|
||||
}
|
||||
h.swap(0, len(h.data)-1)
|
||||
minVal := h.data[len(h.data)-1]
|
||||
h.data = h.data[:len(h.data)-1]
|
||||
h.descend(0)
|
||||
return minVal
|
||||
}
|
||||
|
||||
// Check the top element without removing it.
|
||||
func (h *Heap[T]) Peek() T {
|
||||
if len(h.data) == 0 {
|
||||
panic("heap is empty")
|
||||
}
|
||||
return h.data[0]
|
||||
}
|
||||
|
||||
// Swap child idx with parent idx until the heap property is restored
|
||||
// For example, if we just added a 2 to 0,4,7:
|
||||
// [ 0, 4, 7, 2 ]
|
||||
// Push will call ascend(3) to restore the heap property,
|
||||
// and first we will check idx parent:
|
||||
// i = 3 => parent = (3-1)/2 = 1
|
||||
// 2 < 4 => swap 2 and 4
|
||||
// New i = 1 (current position of 2, previous pos of 4)
|
||||
// so the new parent index for 1 is (1-1)/2 = 0
|
||||
// 2 > 0 => stop (by !h.less(h.data[1], h.data[0]))
|
||||
func (h *Heap[T]) ascend(i int) {
|
||||
for {
|
||||
parent := (i - 1) / 2
|
||||
if i == 0 || !h.less(h.data[i], h.data[parent]) {
|
||||
break
|
||||
}
|
||||
h.swap(i, parent)
|
||||
i = parent
|
||||
}
|
||||
}
|
||||
|
||||
// Swap parent index with the smallest child index until the heap property is restored
|
||||
// For example, let's say we have original slice: [ 0, 2, 6, 4, 3, 7, 9 ]
|
||||
// After popping the root, we swap 0 and 9 and remove 0 (at index 6):
|
||||
// [ 9, 2, 6, 4, 3, 7 ]
|
||||
// Pop will call descend(0) to restore the heap property,
|
||||
// and first we will check index parent:
|
||||
// i = 0 => left = 2*0+1 = 1, right = 2*0+2 = 2
|
||||
// Values lv = 2, rv = 6, rv = 9
|
||||
// 2 < 9 => smallest should be set to left index:
|
||||
// smallest = 1, left = 1, right = 2
|
||||
// 6 > 2, do not change smallest
|
||||
// smallest index (1) != index (0), so we swap 0 and 1 and set index to 1
|
||||
// [ 2, 9, 6, 4, 3, 7 ]
|
||||
// i = 1 => left = 2*1+1 = 3, right = 2*1+2 = 4
|
||||
// Values lv = 4, rv = 3, sv = 9
|
||||
// 4 < 9 => smallest should be set to left index:
|
||||
// smallest = 3, left = 3, right = 4
|
||||
// 3 < 4 => smallest should be set to right index:
|
||||
// smallest = 4, left = 3, right = 4
|
||||
// smallest index (4) != index (1), so we swap 1 and 4 and set index to 4
|
||||
// [ 2, 3, 6, 4, 9, 7 ]
|
||||
// i = 4 => left = 2*4+1 = 9, right = 2*4+2 = 10
|
||||
// meaning there are no children, check that smallest index (4) == index (4) and stop
|
||||
func (h *Heap[T]) descend(i int) {
|
||||
size := len(h.data)
|
||||
for {
|
||||
leftI, rightI := 2*i+1, 2*i+2
|
||||
smallestI := i
|
||||
if leftI < size && h.less(h.data[leftI], h.data[smallestI]) {
|
||||
smallestI = leftI
|
||||
}
|
||||
if rightI < size && h.less(h.data[rightI], h.data[smallestI]) {
|
||||
smallestI = rightI
|
||||
}
|
||||
if smallestI == i {
|
||||
break
|
||||
}
|
||||
h.swap(i, smallestI)
|
||||
i = smallestI
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Heap[T]) swap(i, j int) {
|
||||
h.data[i], h.data[j] = h.data[j], h.data[i]
|
||||
}
|
||||
@@ -1,62 +0,0 @@
|
||||
package aoc
|
||||
|
||||
type PriorityQueue[T comparable] struct {
|
||||
items []T
|
||||
comparator func(a, b T) int
|
||||
}
|
||||
|
||||
func NewPriorityQueue[T comparable](comparator func(a, b T) int) PriorityQueue[T] {
|
||||
return PriorityQueue[T]{comparator: comparator}
|
||||
}
|
||||
func (pq *PriorityQueue[T]) Push(item T) {
|
||||
pq.items = append(pq.items, item)
|
||||
pq.heapifyUp(len(pq.items) - 1)
|
||||
}
|
||||
func (pq *PriorityQueue[T]) Pop() T {
|
||||
top := pq.items[0]
|
||||
lastIndex := len(pq.items) - 1
|
||||
pq.items[0], pq.items[lastIndex] = pq.items[lastIndex], pq.items[0]
|
||||
pq.items = pq.items[:lastIndex]
|
||||
pq.heapifyDown(0)
|
||||
return top
|
||||
}
|
||||
func (pq *PriorityQueue[T]) Peek() T {
|
||||
return pq.items[0]
|
||||
}
|
||||
func (pq *PriorityQueue[T]) Len() int {
|
||||
return len(pq.items)
|
||||
}
|
||||
func (pq *PriorityQueue[T]) IsEmpty() bool {
|
||||
return len(pq.items) == 0
|
||||
}
|
||||
func (pq *PriorityQueue[T]) IsNotEmpty() bool {
|
||||
return !pq.IsEmpty()
|
||||
}
|
||||
func (pq *PriorityQueue[T]) heapifyUp(index int) {
|
||||
for index > 0 {
|
||||
parentIndex := (index - 1) / 2
|
||||
if pq.comparator(pq.items[index], pq.items[parentIndex]) < 0 {
|
||||
pq.items[index], pq.items[parentIndex] = pq.items[parentIndex], pq.items[index]
|
||||
index = parentIndex
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
func (pq *PriorityQueue[T]) heapifyDown(index int) {
|
||||
for {
|
||||
leftChild, rightChild, smallest := 2*index+1, 2*index+2, index
|
||||
if leftChild < len(pq.items) && pq.comparator(pq.items[leftChild], pq.items[smallest]) < 0 {
|
||||
smallest = leftChild
|
||||
}
|
||||
if rightChild < len(pq.items) && pq.comparator(pq.items[rightChild], pq.items[smallest]) < 0 {
|
||||
smallest = rightChild
|
||||
}
|
||||
if smallest != index {
|
||||
pq.items[index], pq.items[smallest] = pq.items[smallest], pq.items[index]
|
||||
index = smallest
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user