快捷導(dǎo)航

Golang實現(xiàn)支持多種類型的set

更新時間：2023年05月25日 08:50:52 作者：nil

在項目開發(fā)中，常常會用到set去重，為什么不寫一個set呢，而且go現(xiàn)在支持了泛型，所以本文就來用Golang實現(xiàn)一個支持多種類型的set呢

寫在前面

今天項目中需要用到string類型的set，想起來項目中不少地方都需要，而且都是用map[string]bool實現(xiàn)的，既然這么多地方用到set去重，為什么不寫一個set呢？而且go現(xiàn)在支持了泛型，為啥不寫一個支持多種類型的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
    }

寫完之后發(fā)現(xiàn)才60多行代碼，非常簡單，而且value使用的是struct{}{}，不占用任何內(nèi)存

單測

    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])
    }