go-study/iterator/set.go

53 lines
794 B
Go
Raw Normal View History

2024-11-12 14:52:48 +08:00
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 func(yield func(T) bool) {
// 迭代有序切片
for _, v := range s.ordered {
if !yield(v) {
return
}
}
}
}