From cb5384671729a02cd6eb9a8fa5a2e44af96c429b Mon Sep 17 00:00:00 2001 From: Brad Fitzpatrick Date: Wed, 5 Jul 2023 19:11:17 -0700 Subject: [PATCH] tempfork/heap: add copy of Go's container/heap but using generics From Go commit 0a48e5cbfabd679e, then with some generics sprinkled about. Updates tailscale/corp#7354 Signed-off-by: Brad Fitzpatrick --- tempfork/heap/heap.go | 121 +++++++++++++++++++++ tempfork/heap/heap_test.go | 216 +++++++++++++++++++++++++++++++++++++ 2 files changed, 337 insertions(+) create mode 100644 tempfork/heap/heap.go create mode 100644 tempfork/heap/heap_test.go diff --git a/tempfork/heap/heap.go b/tempfork/heap/heap.go new file mode 100644 index 000000000..3dfab492a --- /dev/null +++ b/tempfork/heap/heap.go @@ -0,0 +1,121 @@ +// Copyright 2009 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package heap provides heap operations for any type that implements +// heap.Interface. A heap is a tree with the property that each node is the +// minimum-valued node in its subtree. +// +// The minimum element in the tree is the root, at index 0. +// +// A heap is a common way to implement a priority queue. To build a priority +// queue, implement the Heap interface with the (negative) priority as the +// ordering for the Less method, so Push adds items while Pop removes the +// highest-priority item from the queue. The Examples include such an +// implementation; the file example_pq_test.go has the complete source. +// +// This package is a copy of the Go standard library's +// container/heap, but using generics. +package heap + +import "sort" + +// The Interface type describes the requirements +// for a type using the routines in this package. +// Any type that implements it may be used as a +// min-heap with the following invariants (established after +// Init has been called or if the data is empty or sorted): +// +// !h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len() +// +// Note that Push and Pop in this interface are for package heap's +// implementation to call. To add and remove things from the heap, +// use heap.Push and heap.Pop. +type Interface[V any] interface { + sort.Interface + Push(x V) // add x as element Len() + Pop() V // remove and return element Len() - 1. +} + +// Init establishes the heap invariants required by the other routines in this package. +// Init is idempotent with respect to the heap invariants +// and may be called whenever the heap invariants may have been invalidated. +// The complexity is O(n) where n = h.Len(). +func Init[V any](h Interface[V]) { + // heapify + 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[V any](h Interface[V], x V) { + 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(). +// Pop is equivalent to Remove(h, 0). +func Pop[V any](h Interface[V]) V { + n := h.Len() - 1 + h.Swap(0, n) + down(h, 0, n) + 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[V any](h Interface[V], i int) V { + 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[V any](h Interface[V], i int) { + if !down(h, i, h.Len()) { + up(h, i) + } +} + +func up[V any](h Interface[V], j int) { + for { + i := (j - 1) / 2 // parent + if i == j || !h.Less(j, i) { + break + } + h.Swap(i, j) + j = i + } +} + +func down[V any](h Interface[V], i0, n int) bool { + i := i0 + for { + j1 := 2*i + 1 + if j1 >= n || j1 < 0 { // j1 < 0 after int overflow + break + } + j := j1 // left child + if j2 := j1 + 1; j2 < n && h.Less(j2, j1) { + j = j2 // = 2*i + 2 // right child + } + if !h.Less(j, i) { + break + } + h.Swap(i, j) + i = j + } + return i > i0 +} diff --git a/tempfork/heap/heap_test.go b/tempfork/heap/heap_test.go new file mode 100644 index 000000000..c5e96a553 --- /dev/null +++ b/tempfork/heap/heap_test.go @@ -0,0 +1,216 @@ +// Copyright 2009 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package heap + +import ( + "math/rand" + "testing" + + "golang.org/x/exp/constraints" +) + +type myHeap[T constraints.Ordered] []T + +func (h *myHeap[T]) Less(i, j int) bool { + return (*h)[i] < (*h)[j] +} + +func (h *myHeap[T]) Swap(i, j int) { + (*h)[i], (*h)[j] = (*h)[j], (*h)[i] +} + +func (h *myHeap[T]) Len() int { + return len(*h) +} + +func (h *myHeap[T]) Pop() (v T) { + *h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1] + return +} + +func (h *myHeap[T]) Push(v T) { + *h = append(*h, v) +} + +func (h myHeap[T]) verify(t *testing.T, i int) { + t.Helper() + n := h.Len() + j1 := 2*i + 1 + j2 := 2*i + 2 + if j1 < n { + if h.Less(j1, i) { + t.Errorf("heap invariant invalidated [%d] = %v > [%d] = %v", i, h[i], j1, h[j1]) + return + } + h.verify(t, j1) + } + if j2 < n { + if h.Less(j2, i) { + t.Errorf("heap invariant invalidated [%d] = %v > [%d] = %v", i, h[i], j1, h[j2]) + return + } + h.verify(t, j2) + } +} + +func TestInit0(t *testing.T) { + h := new(myHeap[int]) + for i := 20; i > 0; i-- { + h.Push(0) // all elements are the same + } + Init[int](h) + h.verify(t, 0) + + for i := 1; h.Len() > 0; i++ { + x := Pop[int](h) + h.verify(t, 0) + if x != 0 { + t.Errorf("%d.th pop got %d; want %d", i, x, 0) + } + } +} + +func TestInit1(t *testing.T) { + h := new(myHeap[int]) + for i := 20; i > 0; i-- { + h.Push(i) // all elements are different + } + Init[int](h) + h.verify(t, 0) + + for i := 1; h.Len() > 0; i++ { + x := Pop[int](h) + h.verify(t, 0) + if x != i { + t.Errorf("%d.th pop got %d; want %d", i, x, i) + } + } +} + +func Test(t *testing.T) { + h := new(myHeap[int]) + h.verify(t, 0) + + for i := 20; i > 10; i-- { + h.Push(i) + } + Init[int](h) + h.verify(t, 0) + + for i := 10; i > 0; i-- { + Push[int](h, i) + h.verify(t, 0) + } + + for i := 1; h.Len() > 0; i++ { + x := Pop[int](h) + if i < 20 { + Push[int](h, 20+i) + } + h.verify(t, 0) + if x != i { + t.Errorf("%d.th pop got %d; want %d", i, x, i) + } + } +} + +func TestRemove0(t *testing.T) { + h := new(myHeap[int]) + for i := 0; i < 10; i++ { + h.Push(i) + } + h.verify(t, 0) + + for h.Len() > 0 { + i := h.Len() - 1 + x := Remove[int](h, i) + if x != i { + t.Errorf("Remove(%d) got %d; want %d", i, x, i) + } + h.verify(t, 0) + } +} + +func TestRemove1(t *testing.T) { + h := new(myHeap[int]) + for i := 0; i < 10; i++ { + h.Push(i) + } + h.verify(t, 0) + + for i := 0; h.Len() > 0; i++ { + x := Remove[int](h, 0) + if x != i { + t.Errorf("Remove(0) got %d; want %d", x, i) + } + h.verify(t, 0) + } +} + +func TestRemove2(t *testing.T) { + N := 10 + + h := new(myHeap[int]) + for i := 0; i < N; i++ { + h.Push(i) + } + h.verify(t, 0) + + m := make(map[int]bool) + for h.Len() > 0 { + m[Remove[int](h, (h.Len()-1)/2)] = true + h.verify(t, 0) + } + + if len(m) != N { + t.Errorf("len(m) = %d; want %d", len(m), N) + } + for i := 0; i < len(m); i++ { + if !m[i] { + t.Errorf("m[%d] doesn't exist", i) + } + } +} + +func BenchmarkDup(b *testing.B) { + const n = 10000 + h := make(myHeap[int], 0, n) + for i := 0; i < b.N; i++ { + for j := 0; j < n; j++ { + Push[int](&h, 0) // all elements are the same + } + for h.Len() > 0 { + Pop[int](&h) + } + } +} + +func TestFix(t *testing.T) { + h := new(myHeap[int]) + h.verify(t, 0) + + for i := 200; i > 0; i -= 10 { + Push[int](h, i) + } + h.verify(t, 0) + + if (*h)[0] != 10 { + t.Fatalf("Expected head to be 10, was %d", (*h)[0]) + } + (*h)[0] = 210 + Fix[int](h, 0) + h.verify(t, 0) + + for i := 100; i > 0; i-- { + elem := rand.Intn(h.Len()) + if i&1 == 0 { + (*h)[elem] *= 2 + } else { + (*h)[elem] /= 2 + } + Fix[int](h, elem) + h.verify(t, 0) + } +}