Go 中 slice 的 In 功能實(shí)現(xiàn)探索

之前在知乎看到一個(gè)問(wèn)題：為什么 Golang 沒(méi)有像 Python 中 in 一樣的功能？于是，搜了下這個(gè)問(wèn)題，發(fā)現(xiàn)還是有不少人有這樣的疑問(wèn)。

今天來(lái)談?wù)勥@個(gè)話題。

in 是一個(gè)很常用的功能，有些語(yǔ)言中可能也稱(chēng)為 contains，雖然不同語(yǔ)言的表示不同，但基本都是有的。不過(guò)可惜的是，Go 卻沒(méi)有，它即沒(méi)有提供類(lèi)似 Python 操作符 in，也沒(méi)有像其他語(yǔ)言那樣提供這樣的標(biāo)準(zhǔn)庫(kù)函數(shù)，如 PHP 中 in_array。

Go 的哲學(xué)是追求少即是多。我想或許 Go 團(tuán)隊(duì)覺(jué)得這是一個(gè)實(shí)現(xiàn)起來(lái)不足為道的功能吧。

為何說(shuō)微不足道？如果要自己實(shí)現(xiàn)，又該如何做呢？

我所想到的有三種實(shí)現(xiàn)方式，一是遍歷，二是 sort 的二分查找，三是 map 的 key 索引。

本文相關(guān)源碼已經(jīng)上傳在我的 github 上， poloxue/gotin 。

遍歷

遍歷應(yīng)該是我們最容易想到的最簡(jiǎn)單的實(shí)現(xiàn)方式。

示例如下：

func InIntSlice(haystack []int, needle int) bool {
 for _, e := range haystack {
 if e == needle {
 return true
 }
 }

 return false
}

上面演示了如何在一個(gè) []int 類(lèi)型變量中查找指定 int 是否存在的例子，是不是非常簡(jiǎn)單，由此我們也可以感受到我為什么說(shuō)它實(shí)現(xiàn)起來(lái)微不足道。

這個(gè)例子有個(gè)缺陷，它只支持單一類(lèi)型。如果要支持像解釋語(yǔ)言一樣的通用 in 功能，則需借助反射實(shí)現(xiàn)。

代碼如下：

func In(haystack interface{}, needle interface{}) (bool, error) {
 sVal := reflect.ValueOf(haystack)
 kind := sVal.Kind()
 if kind == reflect.Slice || kind == reflect.Array {
 for i := 0; i  sVal.Len(); i++ {
 if sVal.Index(i).Interface() == needle {
 return true, nil
 }
 }

 return false, nil
 }

 return false, ErrUnSupportHaystack
}

為了更加通用，In 函數(shù)的輸入?yún)?shù) haystack 和 needle 都是 interface{} 類(lèi)型。

簡(jiǎn)單說(shuō)說(shuō)輸入?yún)?shù)都是 interface{} 的好處吧，主要有兩點(diǎn)，如下：

其一，haystack 是 interface{} 類(lèi)型，使 in 支持的類(lèi)型不止于 slice，還包括 array。我們看到，函數(shù)內(nèi)部通過(guò)反射對(duì) haystack 進(jìn)行了類(lèi)型檢查，支持 slice（切片）與 array（數(shù)組）。如果是其他類(lèi)型則會(huì)提示錯(cuò)誤，增加新的類(lèi)型支持，如 map，其實(shí)也很簡(jiǎn)單。但不推薦這種方式，因?yàn)橥ㄟ^(guò) _, ok := m[k] 的語(yǔ)法即可達(dá)到 in 的效果。

其二，haystack 是 interface{}，則 []interface{} 也滿足要求，并且 needle 是 interface{}。如此一來(lái)，我們就可以實(shí)現(xiàn)類(lèi)似解釋型語(yǔ)言一樣的效果了。

怎么理解？直接示例說(shuō)明，如下：

gotin.In([]interface{}{1, "two", 3}, "two")

haystack 是 []interface{}{1, "two", 3}，而且 needle 是 interface{}，此時(shí)的值是 "two"。如此看起來(lái)，是不是實(shí)現(xiàn)了解釋型語(yǔ)言中，元素可以是任意類(lèi)型，不必完全相同效果。如此一來(lái)，我們就可以肆意妄為的使用了。

但有一點(diǎn)要說(shuō)明，In 函數(shù)的實(shí)現(xiàn)中有這樣一段代碼：

if sVal.Index(i).Interface() == needle {
 ...
}

Go 中并非任何類(lèi)型都可以使用 == 比較的，如果元素中含有 slice 或 map，則可能會(huì)報(bào)錯(cuò)。

二分查找

以遍歷確認(rèn)元素是否存在有個(gè)缺點(diǎn)，那就是，如果數(shù)組或切片中包含了大量數(shù)據(jù)，比如 1000000 條數(shù)據(jù)，即一百萬(wàn)，最壞的情況是，我們要遍歷 1000000 次才能確認(rèn)，時(shí)間復(fù)雜度 On。

有什么辦法可以降低遍歷次數(shù)？

自然而然地想到的方法是二分查找，它的時(shí)間復(fù)雜度 log2(n) 。但這個(gè)算法有前提，需要依賴有序序列。

于是，第一個(gè)要我們解決的問(wèn)題是使序列有序，Go 的標(biāo)準(zhǔn)庫(kù)已經(jīng)提供了這個(gè)功能，在 sort 包下。

示例代碼如下：

fmt.Println(sort.SortInts([]int{4, 2, 5, 1, 6}))

對(duì)于 []int，我們使用的函數(shù)是 SortInts，如果是其他類(lèi)型切片，sort 也提供了相關(guān)的函數(shù)，比如 []string 可通過(guò) SortStrings 排序。

完成排序就可以進(jìn)行二分查找，幸運(yùn)的是，這個(gè)功能 Go 也提供了，[]int 類(lèi)型對(duì)應(yīng)函數(shù)是 SearchInts。

簡(jiǎn)單介紹下這個(gè)函數(shù)，先看定義：

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

輸入?yún)?shù)容易理解，從切片 a 中搜索 x。重點(diǎn)要說(shuō)下返回值，這對(duì)于我們后面確認(rèn)元素是否存在至關(guān)重要。返回值的含義，返回查找元素在切片中的位置，如果元素不存在，則返回，在保持切片有序情況下，插入該元素應(yīng)該在什么位置。

比如，序列如下：

1 2 6 8 9 11

假設(shè)，x 為 6，查找之后將發(fā)現(xiàn)它的位置在索引 2 處；x 如果是 7，發(fā)現(xiàn)不存在該元素，如果插入序列，將會(huì)放在 6 和 8 之間，索引位置是 3，因而返回值為 3。

代碼測(cè)試下：

fmt.Println(sort.SearchInts([]int{1, 2, 6, 8, 9, 11}, 6)) // 2
fmt.Println(sort.SearchInts([]int{1, 2, 6, 8, 9, 11}, 7)) // 3

如果判斷元素是否在序列中，只要判斷返回位置上的值是否和查找的值相同即可。

但還有另外一種情況，如果插入元素位于序列最后，例如元素值為 12，插入位置即為序列的長(zhǎng)度 6。如果直接查找 6 位置上的元素就可能發(fā)生越界的情況。那怎么辦呢？其實(shí)判斷返回是否大于切片長(zhǎng)度即可，大于則說(shuō)明元素不在切片序列中。

完整的實(shí)現(xiàn)代碼如下：

func SortInIntSlice(haystack []int, needle int) bool {
 sort.Ints(haystack)

 index := sort.SearchInts(haystack, needle)
 return index  len(haystack)  haystack[index] == needle
}

但這還有個(gè)問(wèn)題，對(duì)于無(wú)序的場(chǎng)景，如果每次查詢都要經(jīng)過(guò)一次排序并不劃算。最后能實(shí)現(xiàn)一次排序，稍微修改下代碼。

func InIntSliceSortedFunc(haystack []int) func(int) bool {
 sort.Ints(haystack)

 return func(needle int) bool {
 index := sort.SearchInts(haystack, needle)
 return index  len(haystack)  haystack[index] == needle
 }
}

上面的實(shí)現(xiàn)，我們通過(guò)調(diào)用 InIntSliceSortedFunc 對(duì) haystack 切片排序，并返回一個(gè)可多次使用的函數(shù)。

使用案例如下：

in := gotin.InIntSliceSortedFunc(haystack)

for i := 0; imaxNeedle; i++ {
 if in(i) {
 fmt.Printf("%d is in %v", i, haystack)
 }
}

二分查找的方式有什么不足呢？

我想到的重要一點(diǎn)，要實(shí)現(xiàn)二分查找，元素必須是可排序的，如 int，string，float 類(lèi)型。而對(duì)于結(jié)構(gòu)體、切片、數(shù)組、映射等類(lèi)型，使用起來(lái)就不是那么方便，當(dāng)然，如果要用，也是可以的，不過(guò)需要我們進(jìn)行一些適當(dāng)擴(kuò)展，按指定標(biāo)準(zhǔn)排序，比如結(jié)構(gòu)的某個(gè)成員。

到此，二分查找的 in 實(shí)現(xiàn)就介紹完畢了。

map key

本節(jié)介紹 map key 方式。它的算法復(fù)雜度是 O1，無(wú)論數(shù)據(jù)量多大，查詢性能始終不變。它主要依賴的是 Go 中的 map 數(shù)據(jù)類(lèi)型，通過(guò) hash map 直接檢查 key 是否存在，算法大家應(yīng)該都比較熟悉，通過(guò) key 可直接映射到索引位置。

我們常會(huì)用到這個(gè)方法。

_, ok := m[k]
if ok {
 fmt.Println("Found")
}

那么它和 in 如何結(jié)合呢？一個(gè)案例就說(shuō)明白了這個(gè)問(wèn)題。

假設(shè)，我們有一個(gè) []int 類(lèi)型變量，如下：

s := []int{1, 2, 3}

為了使用 map 的能力檢查某個(gè)元素是否存在，可以將 s 轉(zhuǎn)化 map[int]struct{}。

m := map[interface{}]struct{}{
 1: struct{}{},
 2: struct{}{},
 3: struct{}{},
 4: struct{}{},
}

如果檢查某個(gè)元素是否存在，只需要通過(guò)如下寫(xiě)法即可確定：

k := 4
if _, ok := m[k]; ok {
 fmt.Printf("%d is found\n", k)
}

是不是非常簡(jiǎn)單？

補(bǔ)充一點(diǎn)，關(guān)于這里為什么使用 struct{}，可以閱讀我之前寫(xiě)的一篇關(guān)于Go 中如何使用 set 的文章。

按照這個(gè)思路，實(shí)現(xiàn)函數(shù)如下：

func MapKeyInIntSlice(haystack []int, needle int) bool {
 set := make(map[int]struct{})

 for _ , e := range haystack {
 set[e] = struct{}{}
 }

 _, ok := set[needle]
 return ok
}

實(shí)現(xiàn)起來(lái)不難，但和二分查找有著同樣的問(wèn)題，開(kāi)始要做數(shù)據(jù)處理，將 slice 轉(zhuǎn)化為 map。如果是每次數(shù)據(jù)相同，稍微修改下它的實(shí)現(xiàn)。

func InIntSliceMapKeyFunc(haystack []int) func(int) bool {
 set := make(map[int]struct{})

 for _ , e := range haystack {
 set[e] = struct{}{}
 }

 return func(needle int) bool {
 _, ok := set[needle]
 return ok
 }
}

對(duì)于相同的數(shù)據(jù)，它會(huì)返回一個(gè)可多次使用的 in 函數(shù)，一個(gè)使用案例如下：

in := gotin.InIntSliceMapKeyFunc(haystack)

for i := 0; imaxNeedle; i++ {
 if in(i) {
 fmt.Printf("%d is in %v", i, haystack)
 }
}

對(duì)比前兩種算法，這種方式的處理效率最高，非常適合于大數(shù)據(jù)的處理。接下來(lái)的性能測(cè)試，我們將會(huì)看到效果。

性能

介紹完所有方式，我們來(lái)實(shí)際對(duì)比下每種算法的性能。測(cè)試源碼位于gotin_test.go 文件中。

基準(zhǔn)測(cè)試主要是從數(shù)據(jù)量大小考察不同算法的性能，本文中選擇了三個(gè)量級(jí)的測(cè)試樣本數(shù)據(jù)，分別是 10、1000、1000000。

為便于測(cè)試，首先定義了一個(gè)用于生成 haystack 和 needle 樣本數(shù)據(jù)的函數(shù)。

代碼如下：

func randomHaystackAndNeedle(size int) ([]int, int){
 haystack := make([]int, size)

 for i := 0; isize ; i++{
 haystack[i] = rand.Int()
 }

 return haystack, rand.Int()
}

輸入?yún)?shù)是 size，通過(guò) rand.Int() 隨機(jī)生成切片大小為 size 的 haystack 和 1 個(gè) needle。在基準(zhǔn)測(cè)試用例中，引入這個(gè)隨機(jī)函數(shù)生成數(shù)據(jù)即可。

舉個(gè)例子，如下：

func BenchmarkIn_10(b *testing.B) {
 haystack, needle := randomHaystackAndNeedle(10)

 b.ResetTimer()
 for i := 0; i  b.N; i++ {
 _, _ = gotin.In(haystack, needle)
 }
}

首先，通過(guò) randomHaystackAndNeedle 隨機(jī)生成了一個(gè)含有 10 個(gè)元素的切片。因?yàn)樯蓸颖緮?shù)據(jù)的時(shí)間不應(yīng)該計(jì)入到基準(zhǔn)測(cè)試中，我們使用 b.ResetTimer() 重置了時(shí)間。

其次，壓測(cè)函數(shù)是按照 Test+函數(shù)名+樣本數(shù)據(jù)量 規(guī)則編寫(xiě)，如案例中 BenchmarkIn_10，表示測(cè)試 In 函數(shù)，樣本數(shù)據(jù)量為 10。如果我們要用 1000 數(shù)據(jù)量測(cè)試 InIntSlice，壓測(cè)函數(shù)名為 BenchmarkInIntSlice_1000。

測(cè)試開(kāi)始吧！簡(jiǎn)單說(shuō)下我的筆記本配置，Mac Pro 15 版，16G 內(nèi)存，512 SSD，4 核 8 線程的 CPU。

測(cè)試所有函數(shù)在數(shù)據(jù)量在 10 的情況下的表現(xiàn)。

$ go test -run=none -bench=10$ -benchmem

匹配所有以 10 結(jié)尾的壓測(cè)函數(shù)。

測(cè)試結(jié)果：

goos: darwin
goarch: amd64
pkg: github.com/poloxue/gotin
BenchmarkIn_10-8 3000000 501 ns/op 112 B/op 11 allocs/op
BenchmarkInIntSlice_10-8 200000000 7.47 ns/op 0 B/op 0 allocs/op
BenchmarkInIntSliceSortedFunc_10-8 100000000 22.3 ns/op 0 B/op 0 allocs/op
BenchmarkSortInIntSlice_10-8 10000000 162 ns/op 32 B/op 1 allocs/op
BenchmarkInIntSliceMapKeyFunc_10-8 100000000 17.7 ns/op 0 B/op 0 allocs/op
BenchmarkMapKeyInIntSlice_10-8 3000000 513 ns/op 163 B/op 1 allocs/op
PASS
ok github.com/poloxue/gotin 13.162s

表現(xiàn)最好的并非 SortedFunc 和 MapKeyFunc，而是最簡(jiǎn)單的針對(duì)單類(lèi)型的遍歷查詢，平均耗時(shí) 7.47ns/op，當(dāng)然，另外兩種方式表現(xiàn)也不錯(cuò)，分別是 22.3ns/op 和 17.7ns/op。

表現(xiàn)最差的是 In、SortIn（每次重復(fù)排序） 和 MapKeyIn（每次重復(fù)創(chuàng)建 map）兩種方式，平均耗時(shí)分別為 501ns/op 和 513ns/op。

測(cè)試所有函數(shù)在數(shù)據(jù)量在 1000 的情況下的表現(xiàn)。

$ go test -run=none -bench=1000$ -benchmem

測(cè)試結(jié)果：

goos: darwin
goarch: amd64
pkg: github.com/poloxue/gotin
BenchmarkIn_1000-8 30000 45074 ns/op 8032 B/op 1001 allocs/op
BenchmarkInIntSlice_1000-8 5000000 313 ns/op 0 B/op 0 allocs/op
BenchmarkInIntSliceSortedFunc_1000-8 30000000 44.0 ns/op 0 B/op 0 allocs/op
BenchmarkSortInIntSlice_1000-8 20000 65401 ns/op 32 B/op 1 allocs/op
BenchmarkInIntSliceMapKeyFunc_1000-8 100000000 17.6 ns/op 0 B/op 0 allocs/op
BenchmarkMapKeyInIntSlice_1000-8 20000 82761 ns/op 47798 B/op 65 allocs/op
PASS
ok github.com/poloxue/gotin 11.312s

表現(xiàn)前三依然是 InIntSlice、InIntSliceSortedFunc 和 InIntSliceMapKeyFunc，但這次順序發(fā)生了變化，MapKeyFunc 表現(xiàn)最好，17.6 ns/op，與數(shù)據(jù)量 10 的時(shí)候相比基本無(wú)變化。再次驗(yàn)證了前文的說(shuō)法。

同樣的，數(shù)據(jù)量 1000000 的時(shí)候。

$ go test -run=none -bench=1000000$ -benchmem

測(cè)試結(jié)果如下：

goos: darwin
goarch: amd64
pkg: github.com/poloxue/gotin
BenchmarkIn_1000000-8 30 46099678 ns/op 8000098 B/op 1000001 allocs/op
BenchmarkInIntSlice_1000000-8 3000 424623 ns/op 0 B/op 0 allocs/op
BenchmarkInIntSliceSortedFunc_1000000-8 20000000 72.8 ns/op 0 B/op 0 allocs/op
BenchmarkSortInIntSlice_1000000-8 10 138873420 ns/op 32 B/op 1 allocs/op
BenchmarkInIntSliceMapKeyFunc_1000000-8 100000000 16.5 ns/op 0 B/op 0 allocs/op
BenchmarkMapKeyInIntSlice_1000000-8 10 156215889 ns/op 49824225 B/op 38313 allocs/op
PASS
ok github.com/poloxue/gotin 15.178s

MapKeyFunc 依然表現(xiàn)最好，每次操作用時(shí) 17.2 ns，Sort 次之，而 InIntSlice 呈現(xiàn)線性增加的趨勢(shì)。一般情況下，如果不是對(duì)性能要特殊要求，數(shù)據(jù)量特別大的場(chǎng)景，針對(duì)單類(lèi)型的遍歷已經(jīng)有非常好的性能了。

從測(cè)試結(jié)果可以看出，反射實(shí)現(xiàn)的通用 In 函數(shù)每次執(zhí)行需要進(jìn)行大量的內(nèi)存分配，方便的同時(shí)，也是以犧牲性能為代價(jià)的。

總結(jié)

本文通過(guò)一個(gè)問(wèn)題引出主題，為什么 Go 中沒(méi)有類(lèi)似 Python 的 In 方法。我認(rèn)為，一方面是實(shí)現(xiàn)非常簡(jiǎn)單，沒(méi)有必要。除此以外，另一方面，在不同場(chǎng)景下，我們還需要根據(jù)實(shí)際情況分析用哪種方式實(shí)現(xiàn)，而不是一種固定的方式。

接著，我們介紹了 In 實(shí)現(xiàn)的三種方式，并分析了各自的優(yōu)劣。通過(guò)性能分析測(cè)試，我們能得出大致的結(jié)論，什么方式適合什么場(chǎng)景，但總體還是不能說(shuō)足夠細(xì)致，有興趣的朋友可以繼續(xù)研究下。