利用范型构建 Set

开启掘金成长之旅！这是我参与「掘金日新计划 · 12 月更文挑战」的第12天，点击查看活动详情

你可能对Set这个非常有用的数据集合类型很熟悉。一个集合是一个无序的唯一的集合。通常情况下，集合是使用哈希图实现的，它可以使数值的查询达到O(1)（假设没有哈希图的碰撞）。集合有4个主要操作，使它们特别有用：

• Union (A ⋃ B)：是包含集合A和B中所有元素的集合。
• Intersection (A ∩ B)：是包含在集合A和B中的所有元素的集合。
• Complement (Ac)：是在普遍集合S中但不在A中的元素集合。
• Difference (A − B)：是属于A但不属于B的元素的集合。

让我们开始在Go中定义我们的Set类型。首先，我们要定义什么是集合，通过泛型，我们可以利用约束条件来轻松地扩展集合类型，以处理许多数据类型。

package collections

// Set 用一个 map 来保存唯一的 key
type Set[T comparable] map[T]bool

// NewSet 根据传入的 size 构建一个空的 Set
func NewSet[T comparable](size int) Set[T] {
	return make(Set[T], size)
}

// Add 添加一个 key 到 set 中
func (s Set[T]) Add(key T) {
	s[key] = true
}

// Remove 在 set 中移除一个 key
func (s Set[T]) Remove(key T) {
	delete(s, key)
}

// Contains 检查某个 key 是否在 set 中
func (s Set[T]) Contains(key T) bool {
	return s[key]
}

在第一节中，我们创建了我们的Set类型，它利用了内置的map类型。我们将map的键限制为可比较类型。从文档中，我们知道可比较类型包括：

(booleans, numbers, strings, pointers, channels, arrays of comparable types, structs whose fields are all comparable types)

我们还在我们的类型上添加了一些基本方法，用于添加、删除和检查是否存在。有了这些，我们还没有准备好开始实现我们上面定义的集合操作。让我们从Difference开始：

// Difference 比较两个 set 的不同点
func (a Set[T]) Difference(b Set[T]) Set[T] {
	resultSet := NewSet[T](0)
	for key := range a {
		if !b.Contains(key) {
			resultSet.Add(key)
		}
	}
	return resultSet
}

相当简单的例子。我们简单地创建一个容量为0的新集合（因为我们不知道这个新集合会有多大），然后在集合a上进行迭代，只添加不包含在b中的元素。

接下来的两个操作Union和Intersection遵循类似的模式，但这次我们增加了一个轻微的（或潜在的大）优化。

// Union A union B
func (a Set[T]) Union(b Set[T]) Set[T] {
    small, large := smallLarge(a, b)
    
    for key := range small {
        large.Add(key)
    }
    return large
}

// Intersection A intersect B
func (a Set[T]) Intersection(b Set[T]) Set[T] {
    small, large := smallLarge(a, b)
    
    resultSet := NewSet[T](0)
    for key := range small {
        if large.Contains(key) {
            resultSet.Add(key)
        }
    }
    return resultSet
}

// smallLarge 根据其长度返回小集和大集
func smallLarge[T comparable](a, b Set[T]) (Set[T], Set[T]) {
    small, large := b, a
    if len(b) > len(a) {
        small, large = a, b
    }
    
    return small, large
}

这两种方法都相当简单。在Union中，我们只是在一个集合上进行迭代，将值添加到另一个集合中。在Intersection中，我们要检查a中的值是否也在b中，并返回一个只包含两者元素的集合。

优化来自于在smallLarge(a, b)的调用中区分哪个集合是小的。通过这样做，我们允许循环只迭代两个集合中较小的那个。如果一个集合很大而另一个很小，这有可能节省大量的迭代次数。

然而，在Union中，我们覆盖了大集合，它可能是a或b。如果我们想在Union时保留原始集合，我们必须在两个集合上循环。

我们现在有一个功能齐全的Set包。再做一点工作，我们可以为切片添加helper方法，并添加更多的实用方法，比如检查Equal的方法。

// Equals 比较两个 set 是否相等
func (a Set[T]) Equals(b Set[T]) bool {
	return len(a.Difference(b)) == 0 && len(b.Difference(a)) == 0
}

// ToSlice 将 set 转换为 slice
func (a Set[T]) ToSlice() []T {
	s := make([]T, 0)
	for key := range a {
		s = append(s, key)
	}
	return s
}

// SliceToSet
func SliceToSet[T comparable](s []T) Set[T] {
	set := NewSet[T](len(s))
	for _, item := range s {
		set.Add(item)
	}
	return set
}

// SliceUnion Union 两个 slice
func SliceUnion[T comparable](a, b []T) []T {
	aSet, bSet := SliceToSet(a), SliceToSet(b)
	union := aSet.Union(bSet)
	return union.ToSlice()
}

// SliceIntersection 两个切片的交集。提供的切片不需要是唯一的。不保证顺序。
func SliceIntersection[T comparable](a, b []T) []T {
	aSet, bSet := SliceToSet(a), SliceToSet(b)
	intersection := aSet.Intersection(bSet)
	return intersection.ToSlice()
}

通过上述所有工作，我们能够很轻松执行以下测试代码：

func TestSets(t *testing.T) {
   A := SliceToSet([]int{1, 3, 5})
   B := SliceToSet([]int{0, 1, 2, 3, 4})
  
   union := A.Union(B)
   fmt.Println(union) // map[0:true 1:true 2:true 3:true 4:true 5:true]
  
   C := SliceToSet([]string{"a", "b", "noah"})
   D := SliceToSet([]string{"a", "noah"})
   intersection := C.Intersection(D)
   fmt.Println(intersection) // map[a:true noah:true]
  
   fmt.Println(C.Equals(D)) // false
}

输出：

go test -timeout 30s -run ^TestSets$ go-test/collections -v 
=== RUN TestSets 
map[0:true 1:true 2:true 3:true 4:true 5:true] 
map[a:true noah:true] 
false 
--- PASS: TestSets (0.00s) 
PASS 
ok go-test/collections 0.006s