Merge pull request #1830 from graphitemaster/dale/max_heap

add max heap implementation to slice package

core/slice/heap/heap.odin

@@ -0,0 +1,231 @@
+/*
+	Copyright 2022 Dale Weiler <[email protected]>.
+	Made available under Odin's BSD-3 license.
+
+	List of contributors:
+		Dale Weiler: Initial implementation
+*/
+
+// Package implements a generic max heap in-place on a slice for any type.
+package heap
+
+/*
+	Constructs a max heap in slice given by data with comparator. A max heap is
+	a range of elements which has the following properties:
+
+	1. With N = len(data), for all 0 < i < N, data[(i - 1) / 2] does not compare
+	   less than data[i].
+
+	2. A new element can be added using push in O(log n) time.
+
+	3. The first element can be removed using pop in O(log n) time.
+
+	The comparator compares elements of type T and can be used to construct a
+	max heap (less than) or min heap (greater than) for T.
+*/
+make :: proc(data: []$T, compare: $C) {
+	// amoritize length lookup
+	length := len(data)
+	if length <= 1 do return
+
+	// start from data parent, no need to consider children
+	for start := (length - 2) / 2; start >= 0; start -= 1 {
+		sift_down(data, compare, start)
+	}
+}
+
+/*
+	Inserts the element at the position len(data)-1 into the max heap with
+	comparator.
+
+	At most log(N) comparisons where N = len(data) will be performed.
+*/
+push :: proc(data: []$T, compare: $C) {
+	sift_up(data, compare)
+}
+
+/*
+	Swaps the value in position data[0] and the value in data[len(data)-1] and
+	makes subrange [0, len(data)-1) into a heap. This has the effect of removing
+	the first element from the heap.
+
+	At most 2 * log(N) comparisons where N = len(data) will be performed.
+*/
+pop :: proc(data: []$T, compare: $C) {
+	length := len(data)
+	if length <= 1 do return
+
+	last := length
+
+	// create a hole at 0
+	top := data[0]
+	hole := floyd_sift_down(data, compare)
+	last -= 1
+
+	if hole == last {
+		data[hole] = top
+	} else {
+		data[hole] = data[last]
+		hole += 1
+		data[last] = top
+		sift_up(data[:hole], compare)
+	}
+}
+
+/*
+	Converts the max heap into a sorted range in ascending order. The resulting
+	slice will no longer be a heap after this.
+
+	At most 2 * N * log(N) comparisons where N = len(data) will be performed.
+*/
+sort :: proc(data: []$T, compare: $C) {
+	for n := len(data); n >= 1; n -= 1 {
+		pop(data[:n], compare)
+	}
+}
+
+/*
+	Examines the slice and finds the largest range which is a max-heap. Elements
+	are compared with user-supplied comparison procedure.
+
+	This returns the upper bound of the largest range in the slice which is a
+	max heap. That is, the last index for which data is a max heap.
+
+	At most O(n) comparisons where N = len(data) will be performed.
+*/
+is_heap_until :: proc(data: []$T, compare: $C) -> int {
+	length := len(data)
+	a := 0
+	b := 1
+	for b < length {
+		if compare(data[a], data[b]) {
+			return b
+		}
+		b += 1
+		if b == length || compare(data[a], data[b]) {
+			return b
+		}
+		a += 1
+		b = 2 * a + 1
+	}
+	return length
+}
+
+/*
+	Checks if a given slice is a max heap.
+
+	At most O(n) comparisons where N = len(data) will be performed.
+*/
+is_heap :: #force_inline proc(data: []$T, compare: $C) -> bool {
+	return is_heap_until(data, compare) == len(data)
+}
+
+@(private="file")
+floyd_sift_down :: proc(data: []$T, compare: $C) -> int {
+	length := len(data)
+	assert(length >= 2)
+
+	hole := 0
+	child := 0
+	index := 0
+	for {
+		index += child + 1
+		child = 2 * child + 1
+		if child + 1 < length && compare(data[index], data[index + 1]) {
+			child += 1
+			index += 1
+		}
+
+		data[hole] = data[index]
+		hole = index
+
+		if child > (length - 2) / 2 {
+			return hole
+		}
+	}
+
+	unreachable()
+}
+
+@(private="file")
+sift_down :: proc(data: []$T, compare: $C, start: int) {
+	start := start
+	child := start
+
+	// amoritize length lookup
+	length := len(data)
+
+	// left child of start is at 2 * start + 1
+	// right child of start is at 2 * start + 2
+	if length < 2 || (length - 2) / 2 < child {
+		return
+	}
+
+	child = 2 * child + 1
+
+	if child + 1 < length && compare(data[child], data[child + 1]) {
+		// right child exists and is greater than left child
+		child += 1
+	}
+
+	// check if in heap order
+	if compare(data[child], data[start]) {
+		// start is larger than its largest child
+		return
+	}
+
+	top := data[start]
+	for {
+		// not in heap order, swap parent with its largest child
+		data[start] = data[child]
+		start = child
+
+		if (length - 2) / 2 < child {
+			break
+		}
+
+		// recompute child based off updated parent
+		child = 2 * child + 1
+
+		if child + 1 < length && compare(data[child], data[child + 1]) {
+			// right child exists and is greater than left child
+			child += 1
+		}
+
+		// check if we are in heap order
+		if compare(data[child], top) {
+			break
+		}
+	}
+
+	data[start] = top
+}
+
+@(private="file")
+sift_up :: proc(data: []$T, compare: $C) {
+	// amoritize length lookup
+	length := len(data)
+
+	if length <= 1 do return
+
+	last := length
+	length = (length - 2) / 2
+	index := length
+	last -= 1
+	if compare(data[index], data[last]) {
+		top := data[last]
+		for {
+			data[last] = data[index]
+			last = index
+			if length == 0 {
+				break
+			}
+			length = (length - 1) / 2
+			index = length
+			if !compare(data[index], top) {
+				break
+			}
+		}
+		data[last] = top
+	}
+}