如何用Go判斷元素是否在切片中

2022-06-29 18:02:04

1.問題

如何判斷元素是否在切片中，Golang 並沒有提供直接的庫函數來判斷，最容易想到的實現便是通過遍歷來判斷。

2.遍歷查詢

以字串切片為例，判斷字串切片中是否包含某個字串。

// InSlice 判斷字串是否在 slice 中。
func InSlice(items []string, item string) bool {
	for _, eachItem := range items {
		if eachItem == item {
			return true
		}
	}
	return false
}

這種實現時間複雜度是 O(n)，n 為切片元素個數。

如果切片長度比較短（10以內）或者不是頻繁呼叫，該效能是可以接受的。但是如果切片長度較長且頻繁呼叫，那麼這種方法的效能將無法接受，我們可以藉助 map 優化一波。

3.map 查詢

先將 slice 轉為 map，通過查詢 map 來快速檢視元素是否在 slice 中。

// ConvertStrSlice2Map 將字串 slice 轉為 map[string]struct{}。
func ConvertStrSlice2Map(sl []string) map[string]struct{} {
	set := make(map[string]struct{}, len(sl))
	for _, v := range sl {
		set[v] = struct{}{}
	}
	return set
}

// InMap 判斷字串是否在 map 中。
func InMap(m map[string]struct{}, s string) bool {
	_, ok := m[s]
	return ok
}

注意：使用空結構體 struct{} 作為 value 的型別，因為 struct{} 不佔用任何記憶體空間。

fmt.Println(unsafe.Sizeof(bool(false))) // 1
fmt.Println(unsafe.Sizeof(struct{}{}))  // 0

雖然將 slice 轉為 map 的時間複雜度為 O(n)，但是隻轉換一次可以忽略。查詢元素是否在 map 中的時間複雜度為 O(1)。

4.效能對比

我們可以看下在元素數量為 26 的情況下，取中位元素，做個基準測試（benchmark），對比下二者的查詢效能。

func BenchmarkInSlice(b *testing.B) {
	for i := 0; i < b.N; i++ {
		InSlice(sl, "m")
	}
}

func BenchmarkInMap(b *testing.B) {
	m := ConvertStrSlice2Map(sl)
	for i := 0; i < b.N; i++ {
		InMap(m, "m")
	}
}

執行測試命令輸出：

D:codegotestcontain>go test -bench=.
goos: windows
goarch: amd64
pkg: main/contain
cpu: Intel(R) Core(TM) i7-9700 CPU @ 3.00GHz
BenchmarkInSlice-8 30564058 38.35 ns/op
BenchmarkInMap-8 134556465 8.846 ns/op
PASS
ok main/contain 3.479s

測試結果中，看到函數後面的 -8 個表示執行時對應的 GOMAXPROCS 的值。接著的一串很大的數位表示執行 for 迴圈的次數，也就是呼叫被測試程式碼的次數，最後的38.35 ns/op表示每次需要花費 38.35 納秒。

以上是測試時間預設是 1 秒，也就是1秒的時間，如果想讓測試執行的時間更長，可以通過 -lunchtime 指定，比如 5 秒。

效能對比：

可以預料到的是隨著切片長度增長，效能差距會越來越大。

5.轉換通用化

我們可以藉助空介面 interface{} 來實現任意型別的切片轉換為 map，方便呼叫方使用。

// ToMapSetStrictE converts a slice or array to map set with error strictly.
// The result of map key type is equal to the element type of input.
func ToMapSetStrictE(i interface{}) (interface{}, error) {
	// check param
	if i == nil {
		return nil, fmt.Errorf("unable to converts %#v of type %T to map[interface{}]struct{}", i, i)
	}
	t := reflect.TypeOf(i)
	kind := t.Kind()
	if kind != reflect.Slice && kind != reflect.Array {
		return nil, fmt.Errorf("the input %#v of type %T isn't a slice or array", i, i)
	}
	// execute the convert
	v := reflect.ValueOf(i)
	mT := reflect.MapOf(t.Elem(), reflect.TypeOf(struct{}{}))
	mV := reflect.MakeMapWithSize(mT, v.Len())
	for j := 0; j < v.Len(); j++ {
		mV.SetMapIndex(v.Index(j), reflect.ValueOf(struct{}{}))
	}
	return mV.Interface(), nil
}

func main() {
	var sl = []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z"}
	m, _ := ToMapSetStrictE(sl)
	mSet = m.(map[string]struct{})
	if _, ok := m["m"]; ok {
		fmt.Println("in")
	}
	if _, ok := m["mm"]; !ok {
		fmt.Println("not in")
	}
}

執行輸出：

in
not in

上面的轉換函數ToMapSetStrictE()已經放到開源 Go 工具庫 go-huge-util，可直接通過 go mod 方式 import 使用。

import (
	huge "github.com/dablelv/go-huge-util"
)

// 使用 go-huge-util
m, _ := huge.ToMapSetStrictE(sl)
mSet = m.(map[string]struct{})

// 或使用進一步封裝的函數，不用再斷言
mSet := huge.ToStrMapSetStrict(s)

6.藉助開源庫 golang-set

上面其實是利用 map 實現了一個 set（元素不重複集合），然後再判斷某個 set 中是否存在某個元素。Golang 標準庫並沒有 set，但是我們可以用 map 來間接實現，就像上面那樣子。

如果想使用 set 的完整功能，如初始化、Add、Del、Clear、Contains 等操作，推薦使用 Github 上成熟的開源庫 golang-set，描述中說 Docker 用的也是它。庫中提供了兩種 set 實現，執行緒安全和非執行緒安全的 set。

golang-set 提供了五個生成 set 的函數：

// NewSet creates and returns a reference to an empty set.  Operations
// on the resulting set are thread-safe.
func NewSet(s ...interface{}) Set {}

// NewSetWith creates and returns a new set with the given elements.
// Operations on the resulting set are thread-safe.
func NewSetWith(elts ...interface{}) Set {}

// NewSetFromSlice creates and returns a reference to a set from an
// existing slice.  Operations on the resulting set are thread-safe.
func NewSetFromSlice(s []interface{}) Set {}

// NewThreadUnsafeSet creates and returns a reference to an empty set.
// Operations on the resulting set are not thread-safe.
func NewThreadUnsafeSet() Set {}

// NewThreadUnsafeSetFromSlice creates and returns a reference to a
// set from an existing slice.  Operations on the resulting set are
// not thread-safe.
func NewThreadUnsafeSetFromSlice(s []interface{}) Set {}

下面藉助 golang-set 來判斷切片中是否存在某個元素。

package main

import (
	"fmt"

	mapset "github.com/deckarep/golang-set"
)

func main() {
	var sl = []interface{}{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z"}
	s := mapset.NewSetFromSlice(sl)
	fmt.Println(s.Contains("m"))	// true
	fmt.Println(s.Contains("mm"))	// false
}