slices: rework the APIs of BinarySearch*

For golang/go#50340

Change-Id: If115b2b66d463d5f3788d017924f8dd38867551c
Reviewed-on: https://go-review.googlesource.com/c/exp/+/395414
Reviewed-by: Ian Lance Taylor <[email protected]>
Trust: Eli Bendersky‎ <[email protected]>

Author: eliben

slices/sort.go

@@ -46,25 +46,34 @@ func IsSortedFunc[E any](x []E, less func(a, b E) bool) bool {
 	return true
 }
 
-// BinarySearch searches for target in a sorted slice and returns the smallest
-// index at which target is found. If the target is not found, the index at
-// which it could be inserted into the slice is returned; therefore, if the
-// intention is to find target itself a separate check for equality with the
-// element at the returned index is required.
-func BinarySearch[E constraints.Ordered](x []E, target E) int {
-	return search(len(x), func(i int) bool { return x[i] >= target })
+// BinarySearch searches for target in a sorted slice and returns the position
+// where target is found, or the position where target would appear in the
+// sort order; it also returns a bool saying whether the target is really found
+// in the slice. The slice must be sorted in increasing order.
+func BinarySearch[E constraints.Ordered](x []E, target E) (int, bool) {
+	// search returns the leftmost position where f returns true, or len(x) if f
+	// returns false for all x. This is the insertion position for target in x,
+	// and could point to an element that's either == target or not.
+	pos := search(len(x), func(i int) bool { return x[i] >= target })
+	if pos >= len(x) || x[pos] != target {
+		return pos, false
+	} else {
+		return pos, true
+	}
 }
 
-// BinarySearchFunc uses binary search to find and return the smallest index i
-// in [0, n) at which ok(i) is true, assuming that on the range [0, n),
-// ok(i) == true implies ok(i+1) == true. That is, BinarySearchFunc requires
-// that ok is false for some (possibly empty) prefix of the input range [0, n)
-// and then true for the (possibly empty) remainder; BinarySearchFunc returns
-// the first true index. If there is no such index, BinarySearchFunc returns n.
-// (Note that the "not found" return value is not -1 as in, for instance,
-// strings.Index.) Search calls ok(i) only for i in the range [0, n).
-func BinarySearchFunc[E any](x []E, ok func(E) bool) int {
-	return search(len(x), func(i int) bool { return ok(x[i]) })
+// BinarySearchFunc works like BinarySearch, but uses a custom comparison
+// function. The slice must be sorted in increasing order, where "increasing" is
+// defined by cmp. cmp(a, b) is expected to return an integer comparing the two
+// parameters: 0 if a == b, a negative number if a < b and a positive number if
+// a > b.
+func BinarySearchFunc[E any](x []E, target E, cmp func(E, E) int) (int, bool) {
+	pos := search(len(x), func(i int) bool { return cmp(x[i], target) >= 0 })
+	if pos >= len(x) || cmp(x[pos], target) != 0 {
+		return pos, false
+	} else {
+		return pos, true
+	}
 }
 
 // maxDepth returns a threshold at which quicksort should switch

slices/sort_test.go

@@ -7,6 +7,8 @@ package slices
 import (
 	"math"
 	"math/rand"
+	"strconv"
+	"strings"
 	"testing"
 )
 
@@ -151,31 +153,112 @@ func TestStability(t *testing.T) {
 }
 
 func TestBinarySearch(t *testing.T) {
-	data := []string{"aa", "ad", "ca", "xy"}
+	str1 := []string{"foo"}
+	str2 := []string{"ab", "ca"}
+	str3 := []string{"mo", "qo", "vo"}
+	str4 := []string{"ab", "ad", "ca", "xy"}
+
+	// slice with repeating elements
+	strRepeats := []string{"ba", "ca", "da", "da", "da", "ka", "ma", "ma", "ta"}
+
+	// slice with all element equal
+	strSame := []string{"xx", "xx", "xx"}
+
 	tests := []struct {
-		target string
-		want   int
+		data      []string
+		target    string
+		wantPos   int
+		wantFound bool
 	}{
-		{"aa", 0},
-		{"ab", 1},
-		{"ad", 1},
-		{"ax", 2},
-		{"ca", 2},
-		{"cc", 3},
-		{"dd", 3},
-		{"xy", 3},
-		{"zz", 4},
+		{[]string{}, "foo", 0, false},
+		{[]string{}, "", 0, false},
+
+		{str1, "foo", 0, true},
+		{str1, "bar", 0, false},
+		{str1, "zx", 1, false},
+
+		{str2, "aa", 0, false},
+		{str2, "ab", 0, true},
+		{str2, "ad", 1, false},
+		{str2, "ca", 1, true},
+		{str2, "ra", 2, false},
+
+		{str3, "bb", 0, false},
+		{str3, "mo", 0, true},
+		{str3, "nb", 1, false},
+		{str3, "qo", 1, true},
+		{str3, "tr", 2, false},
+		{str3, "vo", 2, true},
+		{str3, "xr", 3, false},
+
+		{str4, "aa", 0, false},
+		{str4, "ab", 0, true},
+		{str4, "ac", 1, false},
+		{str4, "ad", 1, true},
+		{str4, "ax", 2, false},
+		{str4, "ca", 2, true},
+		{str4, "cc", 3, false},
+		{str4, "dd", 3, false},
+		{str4, "xy", 3, true},
+		{str4, "zz", 4, false},
+
+		{strRepeats, "da", 2, true},
+		{strRepeats, "db", 5, false},
+		{strRepeats, "ma", 6, true},
+		{strRepeats, "mb", 8, false},
+
+		{strSame, "xx", 0, true},
+		{strSame, "ab", 0, false},
+		{strSame, "zz", 3, false},
 	}
 	for _, tt := range tests {
 		t.Run(tt.target, func(t *testing.T) {
-			i := BinarySearch(data, tt.target)
-			if i != tt.want {
-				t.Errorf("BinarySearch want %d, got %d", tt.want, i)
+			{
+				pos, found := BinarySearch(tt.data, tt.target)
+				if pos != tt.wantPos || found != tt.wantFound {
+					t.Errorf("BinarySearch got (%v, %v), want (%v, %v)", pos, found, tt.wantPos, tt.wantFound)
+				}
+			}
+
+			{
+				pos, found := BinarySearchFunc(tt.data, tt.target, strings.Compare)
+				if pos != tt.wantPos || found != tt.wantFound {
+					t.Errorf("BinarySearchFunc got (%v, %v), want (%v, %v)", pos, found, tt.wantPos, tt.wantFound)
+				}
+			}
+		})
+	}
+}
+
+func TestBinarySearchInts(t *testing.T) {
+	data := []int{20, 30, 40, 50, 60, 70, 80, 90}
+	tests := []struct {
+		target    int
+		wantPos   int
+		wantFound bool
+	}{
+		{20, 0, true},
+		{23, 1, false},
+		{43, 3, false},
+		{80, 6, true},
+	}
+	for _, tt := range tests {
+		t.Run(strconv.Itoa(tt.target), func(t *testing.T) {
+			{
+				pos, found := BinarySearch(data, tt.target)
+				if pos != tt.wantPos || found != tt.wantFound {
+					t.Errorf("BinarySearch got (%v, %v), want (%v, %v)", pos, found, tt.wantPos, tt.wantFound)
+				}
 			}
 
-			j := BinarySearchFunc(data, func(s string) bool { return s >= tt.target })
-			if j != tt.want {
-				t.Errorf("BinarySearchFunc want %d, got %d", tt.want, j)
+			{
+				cmp := func(a, b int) int {
+					return a - b
+				}
+				pos, found := BinarySearchFunc(data, tt.target, cmp)
+				if pos != tt.wantPos || found != tt.wantFound {
+					t.Errorf("BinarySearchFunc got (%v, %v), want (%v, %v)", pos, found, tt.wantPos, tt.wantFound)
+				}
 			}
 		})
 	}