sort: binary search for sorted slices

R=r, r2
CC=golang-dev
https://golang.org/cl/2997041

3 files changed, 186 insertions, 0 deletions

diff --git a/src/pkg/sort/Makefile b/src/pkg/sort/Makefile
index 3ea4672409..9deaabfec1 100644
--- a/src/pkg/sort/Makefile
+++ b/src/pkg/sort/Makefile
@@ -6,6 +6,7 @@ include ../../Make.inc
 
 TARG=sort
 GOFILES=\
+	search.go\
 	sort.go\
 
 include ../../Make.pkg
diff --git a/src/pkg/sort/search.go b/src/pkg/sort/search.go
new file mode 100644
index 0000000000..47deda7c9b
--- /dev/null
+++ b/src/pkg/sort/search.go
@@ -0,0 +1,102 @@
+// Copyright 2010 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.
+
+// This file implements binary search.
+
+package sort
+
+// Search uses binary search to find the index i for a value x in an indexable
+// and sorted data structure of n elements.  The argument function f captures
+// the value to be searched for, how the elements are indexed, and how they are
+// sorted.  It will often be passed as a closure.  For instance, given a slice
+// of integers, []data, sorted in ascending order, the function
+//
+//	func(i int) bool { return data[i] <= 23 }
+//
+// can be used to search for the value 23 in data.  The relationship expressed
+// by the function must be "less or equal" if the elements are sorted in ascending
+// order or "greater or equal" if they are sorted in descending order.
+// The function f will be called with values of i in the range 0 to n-1.
+// 
+// For brevity, this discussion assumes ascending sort order. For descending
+// order, replace <= with >=, and swap 'smaller' with 'larger'.
+//
+// If data[0] <= x and x <= data[n-1], Search returns the index i with:
+//
+//	data[i] <= x && x <= data[i+1]
+//
+// where data[n] is assumed to be larger than any x.  Thus, i is the index of x
+// if it is present in the data.  It is the responsibility of the caller to
+// verify the actual presence by testing if data[i] == x.
+//
+// If n == 0 or if x is smaller than any element in data (f is always false),
+// the result is 0.  If x is larger than any element in data (f is always true),
+// the result is n-1.
+//
+// To complete the example above, the following code tries to find the element
+// elem in an integer slice data sorted in ascending order:
+//
+//	elem := 23
+//	i := sort.Search(len(data), func(i int) bool { return data[i] <= elem })
+//	if len(data) > 0 && data[i] == elem {
+//		// elem is present at data[i]
+//	} else {
+//		// elem is not present in data
+//	}
+//
+func Search(n int, f func(int) bool) int {
+	// See "A Method of Programming", E.W. Dijkstra,
+	// for arguments on correctness and efficiency.
+	i, j := 0, n
+	for i+1 < j {
+		h := i + (j-i)/2 // avoid overflow when computing h
+		// i < h < j
+		if f(h) {
+			// data[h] <= x
+			i = h
+		} else {
+			// x < data[h]
+			j = h
+		}
+	}
+	return i
+}
+
+
+// Convenience wrappers for common cases.
+
+// SearchInts searches x in a sorted slice of ints and returns the index
+// as specified by Search. The array must be sorted in ascending order.
+//
+func SearchInts(a []int, x int) int {
+	return Search(len(a), func(i int) bool { return a[i] <= x })
+}
+
+
+// SearchFloats searches x in a sorted slice of floats and returns the index
+// as specified by Search. The array must be sorted in ascending order.
+// 
+func SearchFloats(a []float, x float) int {
+	return Search(len(a), func(i int) bool { return a[i] <= x })
+}
+
+
+// SearchStrings searches x in a sorted slice of strings and returns the index
+// as specified by Search. The array must be sorted in ascending order.
+// 
+func SearchStrings(a []string, x string) int {
+	return Search(len(a), func(i int) bool { return a[i] <= x })
+}
+
+
+// Search returns the result of applying SearchInts to the receiver and x.
+func (p IntArray) Search(x int) int { return SearchInts(p, x) }
+
+
+// Search returns the result of applying SearchFloats to the receiver and x.
+func (p FloatArray) Search(x float) int { return SearchFloats(p, x) }
+
+
+// Search returns the result of applying SearchStrings to the receiver and x.
+func (p StringArray) Search(x string) int { return SearchStrings(p, x) }
diff --git a/src/pkg/sort/search_test.go b/src/pkg/sort/search_test.go
new file mode 100644
index 0000000000..ecd363efbf
--- /dev/null
+++ b/src/pkg/sort/search_test.go
@@ -0,0 +1,83 @@
+// Copyright 2010 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 sort
+
+import "testing"
+
+
+func f(a []int, x int) func(int) bool {
+	return func(i int) bool {
+		return a[i] <= x
+	}
+}
+
+
+var data = []int{0: -10, 1: -5, 2: 0, 3: 1, 4: 2, 5: 3, 6: 5, 7: 7, 8: 11, 9: 100, 10: 100, 11: 100, 12: 1000, 13: 10000}
+
+var tests = []struct {
+	name string
+	n    int
+	f    func(int) bool
+	i    int
+}{
+	{"empty", 0, nil, 0},
+	{"1 1", 1, func(i int) bool { return i <= 1 }, 0},
+	{"1 true", 1, func(i int) bool { return false }, 0},
+	{"1 false", 1, func(i int) bool { return true }, 0},
+	{"1e9 991", 1e9, func(i int) bool { return i <= 991 }, 991},
+	{"1e9 true", 1e9, func(i int) bool { return false }, 0},
+	{"1e9 false", 1e9, func(i int) bool { return true }, 1e9 - 1},
+	{"data -20", len(data), f(data, -20), 0},
+	{"data -10", len(data), f(data, -10), 0},
+	{"data -9", len(data), f(data, -9), 0},
+	{"data -6", len(data), f(data, -6), 0},
+	{"data -5", len(data), f(data, -5), 1},
+	{"data 3", len(data), f(data, 3), 5},
+	{"data 99", len(data), f(data, 99), 8},
+	{"data 100", len(data), f(data, 100), 11},
+	{"data 101", len(data), f(data, 101), 11},
+	{"data 10000", len(data), f(data, 10000), 13},
+	{"data 10001", len(data), f(data, 10001), 13},
+	{"descending a", 7, func(i int) bool { return []int{99, 99, 59, 42, 7, 0, -1, -1}[i] >= 7 }, 4},
+	{"descending 7", 1e9, func(i int) bool { return 1e9-i >= 7 }, 1e9 - 7},
+}
+
+
+func TestSearch(t *testing.T) {
+	for _, e := range tests {
+		i := Search(e.n, e.f)
+		if i != e.i {
+			t.Errorf("%s: expected index %d; got %d", e.name, e.i, i)
+		}
+	}
+}
+
+
+// Smoke tests for convenience wrappers - not comprehensive.
+
+var fdata = []float{0: -3.14, 1: 0, 2: 1, 3: 2, 4: 1000.7}
+var sdata = []string{0: "f", 1: "foo", 2: "foobar", 3: "x"}
+
+var wrappertests = []struct {
+	name   string
+	result int
+	i      int
+}{
+	{"SearchInts", SearchInts(data, 11), 8},
+	{"SearchFloats", SearchFloats(fdata, 2.1), 3},
+	{"SearchStrings", SearchStrings(sdata, ""), 0},
+	{"IntArray.Search", IntArray(data).Search(0), 2},
+	{"FloatArray.Search", FloatArray(fdata).Search(2.0), 3},
+	{"StringArray.Search", StringArray(sdata).Search("x"), 3},
+}
+
+
+func TestSearchWrappers(t *testing.T) {
+	for _, e := range wrappertests {
+		if e.result != e.i {
+			t.Errorf("%s: expected index %d; got %d", e.name, e.i, e.result)
+		}
+	}
+}