go-study/iterator/set.go

package iterator

import (
	"iter"
)

type Set[T comparable] struct {
	ordered []T // 用于保持顺序
	m       map[T]struct{}
}

func NewSet[T comparable]() *Set[T] {
	return &Set[T]{make([]T, 0), make(map[T]struct{})}
}

func (s *Set[T]) Add(e T) {
	if !s.Contains(e) {
		s.ordered = append(s.ordered, e)
	}
	s.m[e] = struct{}{}
}

func (s *Set[T]) Remove(e T) {
	for k, v := range s.ordered {
		if v == e {
			s.ordered = append(s.ordered[:k], s.ordered[k+1:]...)
			break
		}
	}

	delete(s.m, e)
}

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

func (s *Set[T]) Len() int {
	return len(s.m)
}

func (s *Set[T]) All() iter.Seq[T] {
	return s.FilterMap(func(t T) bool {
		return true
	})
}

// FilterMap 筛选迭代
func (s *Set[T]) FilterMap(fn func(T) bool) iter.Seq[T] {
	return func(yield func(T) bool) {
		// 迭代有序切片
		for _, v := range s.ordered {
			if !fn(v) {
				continue
			}

			if !yield(v) {
				return
			}
		}
	}
}