go泛型set

发布日期: 2024-03-13

更新日期: 2024-03-19

文章字数: 452

阅读时长: 2 分

阅读次数:

写在前面

今天项目中需要用到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])
    }