109 lines
1.8 KiB
Go
109 lines
1.8 KiB
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
|
||
}
|
||
}
|
||
}
|
||
}
|