Go 中 set 的实现思考

Go 开发过程中有时我们需要集合（set）这种容器，但 Go 本身未内置这种数据容器，故常常我们需要自己实现，其实实现也很简单。

附，推荐阅读：github.com/Visforest/g...

map[xxx]struct{}

最常用和最容易想到的实现是使用 map，如：

type StrSet struct{
    data map[string]struct{}
}

map 的 value 部分设计为 struct{} 类型是为了节省内存空间。

map[interface{}]struct{}

上面实现的是 string 的 set，如果要其他类型的 set 就得再定义 Int8Set、IntSet、Float32Set 等等，很是繁琐。

很多人可能会选择这样实现 :

type Set struct {
	data map[interface{}]struct{}
}

// New creates a new Set
func New(v ...interface{}) *Set {
	s := &Set{data: map[interface{}]struct{}{}}
	for _, ele := range v {
		s.data[ele] = struct{}{}
	}
	return s
}

// ...

// ToList returns data slice
func (s *Set) ToList() []interface{} {
	var data = make([]interface{}, len(s.data))
	var i int
	for d := range s.data {
		data[i] = d
		i++
	}
	return data
}

这种方式有几个问题：

1. 执行如下代码：

func main() {
	var l1 = []int{1, 2, 3}
	var l2 = []int{4, 5, 6}
	var s = NewSet(l1, l2)
	for _, e := range s.ToList() {
		fmt.Println(e)
	}
}

出错：

panic: runtime error: hash of unhashable type []int

原因很简单，[]int 是不能被 hash 计算的，即不能作为 map 的 key，读者可以查阅 map key允许的类型。interface{} 这种"万金油" 也可能是不合适的。

2. 观察下面代码

func main() {
	var s = NewSet("a", "b", "c")
	var tmp []string
	for _, e := range s.ToList() {
		tmp = append(tmp, e.(string))
	}
	test(tmp)
}

test 函数不能直接拿 s.ToList() 作为入参，必须将 s.ToList() 进行转换为 []string，原因不言自明。

每次都要转换明显损失了编码效率和执行效率。

map[T comparable]struct{}

上面的弊端，可以用 泛型（generics）解决。

定义：

type Set[T comparable] struct {
	data map[T]struct{}
}

// New creates a new Set
func NewSet[T comparable](v ...T) *Set[T] {
	s := &Set[T]{data: map[T]struct{}{}}
	for _, ele := range v {
		s.data[ele] = struct{}{}
	}
	return s
}

func (s *Set[T]) Add(v ...T) {
	for _, ele := range v {
		s.data[ele] = struct{}{}
	}
}

// ...

// ToList returns data slice
func (s *Set[T]) ToList() []T {
	var data = make([]T, len(s.data))
	var i int
	for d := range s.data {
		data[i] = d
		i++
	}
	return data
}

使用：

func test1(data []string) {
	// ...
}

func test2(data []float64) {
	// ...
}

func main() {
	var s1 = NewSet("a", "b", "c")
	test1(s1.ToList())

	var s2 = NewSet(1.3, 2.2, 3)
	test2(s2.ToList())
}

type IntSet = Set[int]

上面的 Set 是个通用 set，类型混用时自己可能会被误导。我们可以定义专用数据类型的 set，且代码不需要很多。

type IntSet = Set[int]

func NewIntSet(v ...int) *IntSet {
	return NewSet[int](v...)
}

使用：

func main() {
	var s = NewIntSet(1, 2, 3)
	test3(s.ToList())

	// 编译错误
	// s.Add("a", "b", "c")
}

fifo set

通常 set 是无序的，上面的实现也都是无序的，但有的场景下我们需要有序的 set，比如fifo set，sorted set。这里以 fifo set 为例，讨论下其实现。

为了兼顾查找效率和有序特性，可以使用 map + array / double linkedlist，考虑到数据的添加、删除以及内存使用，double linkedlist 有比 array 显著的优势。

type setNode[T comparable] struct {
	val  T
	pre  *setNode[T]
	next *setNode[T]
}

type FifoSet[T comparable] struct {
	head *setNode[T]
	tail *setNode[T]
	data map[T]*setNode[T]
}

// add data, make it first in first out
func (l *FifoSet[T]) Add(v ...T) {
	if len(v) == 0 {
		return
	}

	var i int
	if l.head == nil {
		// first node
		n := &setNode[T]{
			val: v[i],
		}
		l.head = n
		l.tail = n
		l.data[v[i]] = n
		i++
	}
	for ; i < len(v); i++ {
		if _, ok := l.data[v[i]]; !ok {
            // when missing, insert
			n := &setNode[T]{
				val:  v[i],
				pre:  l.tail,
				next: nil,
			}
			l.tail.next = n
			l.tail = n
			l.data[v[i]] = n
		}
	}
}

使用：

func main() {
	var s = NewFifoSet[string]()
	s.Add("e", "a", "b", "a", "c", "b")
	// e
	// a
    // b
	// c
	for _, v := range s.ToList() {
		fmt.Println(v)
	}
}

sorted set

其实 sorted set 与 fifo set 实现很像，只是略有区别，这里就略过了。

有兴趣的可以阅读笔者的 github.com/Visforest/g...，或者自己尝试自己实现下。

欢迎交流。