技术

Go二分查找sort.Search

Go语言, 算法

之前对Go sort包的印象一直是只能做排序，毕竟包名就叫sort嘛。后来在一次刷leetcode的时候，发现官方题解做二分搜索的时候用了sort包里的Search函数，惊讶sort包还封装了二分查找功能。于是看了下sort.Search的源码，发现源码其实也写的很简单，就是一个普通的二分：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
func Search(n int, f func(int) bool) int {
	// Define f(-1) == false and f(n) == true.
	// Invariant: f(i-1) == false, f(j) == true.
	i, j := 0, n
	for i < j {
		h := int(uint(i+j) >> 1) // avoid overflow when computing h
		// i ≤ h < j
		if !f(h) {
			i = h + 1 // preserves f(i-1) == false
		} else {
			j = h // preserves f(j) == true
		}
	}
	// i == j, f(i-1) == false, and f(j) (= f(i)) == true  =>  answer is i.
	return i
}

sort.Search的作用就是通过二分法，找到满足条件（即函数f返回值为true）的最小的索引。
根据sort.Search的这个作用，Go官方的包还给出了几个函数。比如sort.SearchInts，其函数签名为：

1
func SearchInts(a []int, x int) int

sort.SearchInts中调用了sort.Search，传入的函数f是：

1
2
3
func(i int) bool { 
    return a[i] >= x 
}

看到这个函数f，其实sort.SearchInts的作用就很明显了，即从传入的切片a中找到第一个大于等于x的索引。这个其实就是C++中lower_bound的功能。

sort.SearchInts已经封装了在一个已正向排序的切片中进行lower_bound查询的功能。如果要进行其他类型的二分查询，修改以传入到sort.Search的第二个参数f即可。这里举几个常用的例子：

查找第一个大于x的索引，即upper_bound

1
2
3
4
5
func UpperBound(a []int, x int) int {
    return sort.Search(len(a), func(i int) bool {
        return a[i] > x
    })
}

查询第一个小于x的索引

1
2
3
4
5
func SearchLess(a []int, x int) int {
    return sort.Search(len(a), func(i int) bool {
        return a[i] < x
    })
}

查询第一个小于等于x的索引

1
2
3
4
5
func SearchLessOrEqual(a []int, x int) int {
    return sort.Search(len(a), func(i int) bool {
        return a[i] <= x
    })
}

当然，传入的切片a必须是保证有序的。

分享到