From 3aed2695a3edfba7c5d95ab86b8854cc1b6dc6a8 Mon Sep 17 00:00:00 2001 From: Luke Shumaker Date: Fri, 15 Jul 2022 20:27:50 -0600 Subject: add binary search tools to lib/slices --- lib/slices/sliceutil.go | 107 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) (limited to 'lib/slices') diff --git a/lib/slices/sliceutil.go b/lib/slices/sliceutil.go index faaffcb..f95c7e7 100644 --- a/lib/slices/sliceutil.go +++ b/lib/slices/sliceutil.go @@ -73,3 +73,110 @@ func Sort[T constraints.Ordered](slice []T) { return slice[i] < slice[j] }) } + +// returns (a+b)/2, but avoids overflow +func avg(a, b int) int { + return int(uint(a+b) >> 1) +} + +func min(a, b int) int { + if a < b { + return a + } + return b +} + +func max(a, b int) int { + if a > b { + return a + } + return b +} + +// Search the slice for a value for which `fn(slice[i]) = 0`. +// +// + + + 0 0 0 - - - +// ^ ^ ^ +// any of +// +// You can conceptualize `fn` as subtraction: +// +// func(straw T) int { +// return needle - straw +// } +func Search[T any](slice []T, fn func(T) int) (int, bool) { + beg, end := 0, len(slice) + for beg < end { + midpoint := avg(beg, end) + direction := fn(slice[midpoint]) + switch { + case direction < 0: + end = midpoint + case direction > 0: + beg = midpoint + 1 + case direction == 0: + return midpoint, true + } + } + return 0, false +} + +// Search the slice for the left-most value for which `fn(slice[i]) = 0`. +// +// + + + 0 0 0 - - - +// ^ +// +// You can conceptualize `fn` as subtraction: +// +// func(straw T) int { +// return needle - straw +// } +func SearchLowest[T any](slice []T, fn func(T) int) (int, bool) { + lastBad, firstGood, firstBad := -1, len(slice), len(slice) + for lastBad+1 < min(firstGood, firstBad) { + midpoint := avg(lastBad, min(firstGood, firstBad)) + direction := fn(slice[midpoint]) + switch { + case direction < 0: + firstBad = midpoint + case direction > 0: + lastBad = midpoint + default: + firstGood = midpoint + } + } + if firstGood == len(slice) { + return 0, false + } + return firstGood, true +} + +// Search the slice for the right-most value for which `fn(slice[i]) = 0`. +// +// + + + 0 0 0 - - - +// ^ +// +// You can conceptualize `fn` as subtraction: +// +// func(straw T) int { +// return needle - straw +// } +func SearchHighest[T any](slice []T, fn func(T) int) (int, bool) { + lastBad, lastGood, firstBad := -1, -1, len(slice) + for max(lastBad, lastGood)+1 < firstBad { + midpoint := avg(max(lastBad, lastGood), firstBad) + direction := fn(slice[midpoint]) + switch { + case direction < 0: + firstBad = midpoint + case direction > 0: + lastBad = midpoint + default: + lastGood = midpoint + } + } + if lastGood < 0 { + return 0, false + } + return lastGood, true +} -- cgit v1.2.3-2-g168b