go-study/iterator/set.go

package iterator

import (
	"iter"
	"slices"
)

// Set 集合
type Set[T comparable] struct {
	m map[T]struct{}
}

// OrderedSet 有序集合
type OrderedSet[T comparable] struct {
	ordered []T // 用于保持顺序
	Set[T]
}

func NewSet[T comparable]() Set[T] {
	return Set[T]{make(map[T]struct{})}
}

func NewOrderedSet[T comparable]() *OrderedSet[T] {
	return &OrderedSet[T]{make([]T, 0), NewSet[T]()}
}

func (s Set[T]) FromSlice(slice []T) {
	for v := range slices.Values(slice) {
		s.Add(v)
	}
}

func (s *OrderedSet[T]) FromSlice(slice []T) {
	for v := range slices.Values(slice) {
		s.Add(v)
	}
}

func (s Set[T]) Add(e T) {
	s.m[e] = struct{}{}
}

func (s *OrderedSet[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) {
	delete(s.m, e)
}

func (s *OrderedSet[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) bool {
		return true
	})
}

// All 有序集合的推迭代器，该方法不能省略，否则会走到子集合的FilterMap方法中去
func (s *OrderedSet[T]) All() iter.Seq[T] {
	return s.FilterMap(func(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.m {
			if fn(v) && !yield(v) {
				return
			}
		}
	}
}

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