From f6cff4143246f79ed480bfb1154b9adbf89f9179 Mon Sep 17 00:00:00 2001 From: Matthew Kim <38759997+friendlymatthew@users.noreply.github.com> Date: Thu, 13 Jun 2024 11:08:08 -0400 Subject: [PATCH] added minmax --- pkg/hnsw/minmax.go | 192 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 192 insertions(+) create mode 100644 pkg/hnsw/minmax.go diff --git a/pkg/hnsw/minmax.go b/pkg/hnsw/minmax.go new file mode 100644 index 0000000..b6658c1 --- /dev/null +++ b/pkg/hnsw/minmax.go @@ -0,0 +1,192 @@ +// Package minmaxheap provides min-max heap operations for any type that +// implements heap.Interface. A min-max heap can be used to implement a +// double-ended priority queue. +// +// Min-max heap implementation from the 1986 paper "Min-Max Heaps and +// Generalized Priority Queues" by Atkinson et. al. +// https://doi.org/10.1145/6617.6621. + +package hnsw + +import ( + "container/heap" + "math/bits" +) + +// Interface copied from the heap package, so code that imports minmaxheap does +// not also have to import "container/heap". +type Interface = heap.Interface + +func level(i int) int { + // floor(log2(i + 1)) + return bits.Len(uint(i)+1) - 1 +} + +func isMinLevel(i int) bool { + return level(i)%2 == 0 +} + +func lchild(i int) int { + return i*2 + 1 +} + +func rchild(i int) int { + return i*2 + 2 +} + +func parent(i int) int { + return (i - 1) / 2 +} + +func hasParent(i int) bool { + return i > 0 +} + +func hasGrandparent(i int) bool { + return i > 2 +} + +func grandparent(i int) int { + return parent(parent(i)) +} + +func down(h Interface, i, n int) bool { + min := isMinLevel(i) + i0 := i + for { + m := i + + l := lchild(i) + if l >= n || l < 0 /* overflow */ { + break + } + if h.Less(l, m) == min { + m = l + } + + r := rchild(i) + if r < n && h.Less(r, m) == min { + m = r + } + + // grandchildren are contiguous i*4+3+{0,1,2,3} + for g := lchild(l); g < n && g <= rchild(r); g++ { + if h.Less(g, m) == min { + m = g + } + } + + if m == i { + break + } + + h.Swap(i, m) + + if m == l || m == r { + break + } + + // m is grandchild + p := parent(m) + if h.Less(p, m) == min { + h.Swap(m, p) + } + i = m + } + return i > i0 +} + +func up(h Interface, i int) { + min := isMinLevel(i) + + if hasParent(i) { + p := parent(i) + if h.Less(p, i) == min { + h.Swap(i, p) + min = !min + i = p + } + } + + for hasGrandparent(i) { + g := grandparent(i) + if h.Less(i, g) != min { + return + } + + h.Swap(i, g) + i = g + } +} + +// Init establishes the heap invariants required by the other routines in this +// package. Init may be called whenever the heap invariants may have been +// invalidated. +// The complexity is O(n) where n = h.Len(). +func Init(h Interface) { + n := h.Len() + for i := n/2 - 1; i >= 0; i-- { + down(h, i, n) + } +} + +// Push pushes the element x onto the heap. +// The complexity is O(log n) where n = h.Len(). +func Push(h Interface, x interface{}) { + h.Push(x) + up(h, h.Len()-1) +} + +// Pop removes and returns the minimum element (according to Less) from the heap. +// The complexity is O(log n) where n = h.Len(). +func Pop(h Interface) interface{} { + n := h.Len() - 1 + h.Swap(0, n) + down(h, 0, n) + return h.Pop() +} + +// PopMax removes and returns the maximum element (according to Less) from the heap. +// The complexity is O(log n) where n = h.Len(). +func PopMax(h Interface) interface{} { + n := h.Len() + + i := 0 + l := lchild(0) + if l < n && !h.Less(l, i) { + i = l + } + + r := rchild(0) + if r < n && !h.Less(r, i) { + i = r + } + + h.Swap(i, n-1) + down(h, i, n-1) + return h.Pop() +} + +// Remove removes and returns the element at index i from the heap. +// The complexity is O(log n) where n = h.Len(). +func Remove(h Interface, i int) interface{} { + n := h.Len() - 1 + if n != i { + h.Swap(i, n) + if !down(h, i, n) { + up(h, i) + } + } + return h.Pop() +} + +// Fix re-establishes the heap ordering after the element at index i has +// changed its value. Changing the value of the element at index i and then +// calling Fix is equivalent to, but less expensive than, calling Remove(h, i) +// followed by a Push of the new value. +// The complexity is O(log n) where n = h.Len(). +func Fix(h Interface, i int) { + if !down(h, i, h.Len()) { + up(h, i) + } +}