go泛型set

506 阅读1分钟

写在前面

今天项目中需要用到string类型的set,想起来项目中不少地方都需要,而且都是用map[string]bool实现的,既然这么多地方用到set去重,为什么不写一个set呢?而且go现在支持了泛型,为啥不写一个支持多种类型的set呢?说干就干

代码

    package set

    var (
    	ElemValue = struct{}{}
    )

    type SetType interface {
    	int | int32 | int64 | string | float32 | float64
    }

    // Set[T SetType]. 集合,线程不安全
    type Set[T SetType] interface {
    	Add(T)
    	Remove(T)
    	Contains(T) bool
    	Empty() bool
    	Values() []T
    }

    type set[T SetType] struct {
    	m map[T]struct{}
    }

    func NewSet[T SetType](ss ...T) Set[T] {
    	newSet := set[T]{
    		m: make(map[T]struct{}, len(ss)),
    	}

    	for _, s := range ss {
    		newSet.Add(s)
    	}

    	return &newSet
    }

    func (s *set[T]) Add(elem T) {
    	s.m[elem] = ElemValue
    }

    func (s *set[T]) Remove(elem T) {
    	delete(s.m, elem)
    }

    func (s *set[T]) Contains(elem T) bool {
    	_, ok := s.m[elem]
    	return ok
    }

    func (s *set[T]) Empty() bool {
    	empty := true
    	for _, _ = range s.m {
    		empty = false
    		break
    	}
    	return empty
    }

    func (s *set[T]) Values() []T {
    	ss := make([]T, 0)
    	for k, _ := range s.m {
    		ss = append(ss, k)
    	}

    	return ss
    }

写完之后发现才60多行代码,非常简单,而且value使用的是struct{}{},不占用任何内存

单测

    package set_test

    import (
    	"set"
    	"testing"

    	"github.com/stretchr/testify/assert"
    )

    func TestStringSet(t *testing.T) {
    	newSet := set.NewSet[string]("a", "b")
    	assert.False(t, newSet.Empty())

    	newSet.Add("c")
    	newSet.Remove("a")

    	assert.False(t, newSet.Contains("a"))
    	assert.True(t, newSet.Contains("c"))

    	newSet.Remove("c")

    	result := newSet.Values()
    	assert.Equal(t, "b", result[0])
    }

    func TestInt64Set(t *testing.T) {
    	newSet := set.NewSet[int64](1, 2)
    	assert.False(t, newSet.Empty())

    	newSet.Add(3)
    	newSet.Remove(1)

    	assert.False(t, newSet.Contains(1))
    	assert.True(t, newSet.Contains(3))

    	newSet.Remove(3)

    	result := newSet.Values()
    	assert.Equal(t, int64(2), result[0])
    }

造轮子感想

之前没有用过go 泛型,今天写这个set体验了一下泛型,确实挺好用的。本人喜欢造轮子、支持造轮子。只有这样才能提升自己,一起学习吧

源代码:github.com/ZBIGBEAR/se…