lib/algo/legacy/heaps.viv

# Generated by tools/algo/legacy.py from the corpus: do not edit.
# The two heaps of the TheAlgorithms ports: heap.py's max-heap of numbers, and the (cost, vertex) min-heap Dijkstra uses (heapq.heappush and heappop).
# Element comparisons are routed through lib/algo/legacy/order.viv; nothing else changes.
from algo.legacy.order import lg_gt, lg_lt, lg_ne

# from study/idioms/thealgorithms-python/structures.viv: heap_new heap_swap max_heapify build_max_heap extract_max heap_insert heap_sort_in_place heap_items
H_SIZE = 0.0
H_ITEMS = 1.0

def heap_new(capacity) -> Arr {
    return zeros(H_ITEMS + capacity)
}

def heap_swap(h: Arr, i, j) {
    t = h[H_ITEMS + i]
    h[H_ITEMS + i] = h[H_ITEMS + j]
    h[H_ITEMS + j] = t
}

#| Heap.max_heapify. Python's left_child_idx returns None past the end; here `and`
#| checks the bound before the comparison reads the child.
def max_heapify(h: Arr, index) {
    if index < h[H_SIZE] {
        violation = index
        left = 2.0 * index + 1.0
        right = 2.0 * index + 2.0
        if left < h[H_SIZE] and lg_gt(h[H_ITEMS + left], h[H_ITEMS + violation]) {
            violation = left
        }
        if right < h[H_SIZE] and lg_gt(h[H_ITEMS + right], h[H_ITEMS + violation]) {
            violation = right
        }
        if violation != index {
            heap_swap(h, violation, index)
            max_heapify(h, violation)
        }
    }
}

def build_max_heap(h: Arr, collection: Arr) {
    for i to len(collection) {
        h[H_ITEMS + i] = collection[i]
    }
    h[H_SIZE] = len(collection)
    if h[H_SIZE] > 1.0 {
        for i to h[H_SIZE] // 2.0 reversed {
            max_heapify(h, i)
        }
    }
}

#| Heap.extract_max. The caller checks the size first: Python raises on an empty heap.
def extract_max(h: Arr) {
    me = h[H_ITEMS]
    last = h[H_ITEMS + h[H_SIZE] - 1.0]
    h[H_SIZE] -= 1.0
    if h[H_SIZE] >= 1.0 {
        h[H_ITEMS] = last
        max_heapify(h, 0.0)
    }
    return me
}

def heap_insert(h: Arr, value) {
    h[H_ITEMS + h[H_SIZE]] = value
    idx = (h[H_SIZE] - 1.0) // 2.0
    h[H_SIZE] += 1.0
    while idx >= 0.0 {
        max_heapify(h, idx)
        idx = (idx - 1.0) // 2.0
    }
}

def heap_sort_in_place(h: Arr) {
    size = h[H_SIZE]
    for j from 1.0 to size reversed {
        heap_swap(h, 0.0, j)
        h[H_SIZE] -= 1.0
        max_heapify(h, 0.0)
    }
    h[H_SIZE] = size
}

#| The items as a list (Python `repr(heap)`).
def heap_items(h: Arr) -> Arr {
    return h[H_ITEMS:H_ITEMS + h[H_SIZE]]
}

# from study/idioms/thealgorithms-python/graphs.viv: less swap_rows heappush heappop

#| A binary min-heap of (cost, vertex) rows, ordered like Python tuples, in the first
#| `size` rows of h. heapq.heappush: append, then sift the new row up.
def less(h: Arr, i, j) -> Bool {
    if lg_ne(h[i, 0], h[j, 0]) {
        return lg_lt(h[i, 0], h[j, 0])
    }
    return lg_lt(h[i, 1], h[j, 1])
}

def swap_rows(h: Arr, i, j) {
    for c to 2 {
        t = h[i, c]
        h[i, c] = h[j, c]
        h[j, c] = t
    }
}

def heappush(h: Arr, size, cost, vertex) {
    h[size, 0] = cost
    h[size, 1] = vertex
    i = size
    while i > 0.0 {
        parent = floor((i - 1.0) / 2.0)
        if not less(h, i, parent) {
            break
        }
        swap_rows(h, i, parent)
        i = parent
    }
    return size + 1.0
}

#| heapq.heappop, arranged the way C++'s pop_heap is: the smallest row is moved to
#| row size - 1, the rest is re-heaped, and the new size (size - 1) is returned, so the
#| caller reads the popped pair at h[size - 1, ...].
def heappop(h: Arr, size) {
    last = size - 1.0
    swap_rows(h, 0.0, last)
    i = 0.0
    while true {
        child = 2.0 * i + 1.0
        if child >= last {
            break
        }
        if child + 1.0 < last and less(h, child + 1.0, child) {
            child += 1.0
        }
        if not less(h, child, i) {
            break
        }
        swap_rows(h, i, child)
        i = child
    }
    return last
}