tiny-rdm/backend/utils/coll/set.go

package coll

import (
	"fmt"
	json "github.com/bytedance/sonic"
	"sort"
	. "tinyrdm/backend/utils"
	"tinyrdm/backend/utils/rand"
)

type Void struct{}

// Set 集合, 存放不重复的元素
type Set[T Hashable] map[T]Void

// type Set[T Hashable] struct {
//	data map[T]Void
// }

func NewSet[T Hashable](elems ...T) Set[T] {
	if len(elems) > 0 {
		data := make(Set[T], len(elems))
		for _, e := range elems {
			data[e] = Void{}
		}
		return data
	} else {
		return Set[T]{}
	}
}

// Add 添加元素
func (s Set[T]) Add(elem T) bool {
	if s == nil {
		return false
	}
	if _, exists := s[elem]; !exists {
		s[elem] = Void{}
		return true
	}
	return false
}

// AddN 添加多个元素
func (s Set[T]) AddN(elems ...T) int {
	if s == nil {
		return 0
	}
	addCount := 0
	var exists bool
	for _, elem := range elems {
		if _, exists = s[elem]; !exists {
			s[elem] = Void{}
			addCount += 1
		}
	}
	return addCount
}

// Merge 合并其他集合
func (s Set[T]) Merge(other Set[T]) int {
	return s.AddN(other.ToSlice()...)
}

// Contains 判断是否存在指定元素
func (s Set[T]) Contains(elem T) bool {
	if s == nil {
		return false
	}
	_, exists := s[elem]
	return exists
}

// ContainAny 判断是否包含任意元素
func (s Set[T]) ContainAny(elems ...T) bool {
	if s == nil {
		return false
	}
	var exists bool
	for _, elem := range elems {
		if _, exists = s[elem]; exists {
			return true
		}
	}
	return false
}

// Equals 判断两个集合内元素是否一致
func (s Set[T]) Equals(other Set[T]) bool {
	if s.Size() != other.Size() {
		return false
	}
	for elem := range s {
		if !other.Contains(elem) {
			return false
		}
	}
	return true
}

// ContainAll 判断是否包含所有元素
func (s Set[T]) ContainAll(elems ...T) bool {
	if s == nil {
		return false
	}
	var exists bool
	for _, elem := range elems {
		if _, exists = s[elem]; !exists {
			return false
		}
	}
	return true
}

// Remove 移除元素
func (s Set[T]) Remove(elem T) bool {
	if s == nil {
		return false
	}
	if _, exists := s[elem]; exists {
		delete(s, elem)
		return true
	}
	return false
}

// RemoveN 移除多个元素
func (s Set[T]) RemoveN(elems ...T) int {
	if s == nil {
		return 0
	}
	var exists bool
	removeCnt := 0
	for _, elem := range elems {
		if _, exists = s[elem]; exists {
			delete(s, elem)
			removeCnt += 1
		}
	}
	return removeCnt
}

// RemoveSub 移除子集
func (s Set[T]) RemoveSub(subSet Set[T]) int {
	if s == nil {
		return 0
	}
	var exists bool
	removeCnt := 0
	for elem := range subSet {
		if _, exists = s[elem]; exists {
			delete(s, elem)
			removeCnt += 1
		}
	}
	return removeCnt
}

// Filter 根据条件筛出符合的元素
func (s Set[T]) Filter(filterFunc func(i T) bool) []T {
	ret := []T{}
	for v := range s {
		if filterFunc(v) {
			ret = append(ret, v)
		}
	}
	return ret
}

// RandomElem 随机抽取一个元素
// @param remove 随机出来的元素是否同时从集合中移除
// @return 抽取的元素
// @return 是否抽取成功
func (s Set[T]) RandomElem(remove bool) (T, bool) {
	size := s.Size()
	if size > 0 {
		selIdx := rand.Intn(size)
		idx := 0
		for elem := range s {
			if idx == selIdx {
				if remove {
					delete(s, elem)
				}
				return elem, true
			} else {
				idx++
			}
		}
	}

	var r T
	return r, false
}

// Size 集合长度
func (s Set[T]) Size() int {
	return len(s)
}

// IsEmpty 判断是否为空
func (s Set[T]) IsEmpty() bool {
	return len(s) <= 0
}

// Clear 清空集合
func (s Set[T]) Clear() {
	for elem := range s {
		delete(s, elem)
	}
}

// ToSlice 转为切片
func (s Set[T]) ToSlice() []T {
	size := len(s)
	if size <= 0 {
		return []T{}
	}

	ret := make([]T, 0, size)
	for elem := range s {
		ret = append(ret, elem)
	}
	return ret
}

// ToSortedSlice 转为排序好的切片
func (s Set[T]) ToSortedSlice(sortFunc func(v1, v2 T) bool) []T {
	list := s.ToSlice()
	sort.Slice(list, func(i, j int) bool {
		return sortFunc(list[i], list[j])
	})
	return list
}

// Each 遍历检索每个元素
func (s Set[T]) Each(eachFunc func(T)) {
	if len(s) <= 0 {
		return
	}
	for elem := range s {
		eachFunc(elem)
	}
}

// Clone 克隆
func (s Set[T]) Clone() Set[T] {
	if s == nil {
		return nil
	}

	other := NewSet[T]()
	for elem := range s {
		other[elem] = Void{}
	}
	return other
}

func (s Set[T]) String() string {
	arr := s.ToSlice()
	return fmt.Sprintf("%v", arr)
}

// MarshalJSON to output non base64 encoded []byte
func (s Set[T]) MarshalJSON() ([]byte, error) {
	if s == nil {
		return []byte("null"), nil
	}
	t := s.ToSlice()
	return json.Marshal(t)
}

// UnmarshalJSON to deserialize []byte
func (s *Set[T]) UnmarshalJSON(b []byte) error {
	t := []T{}
	err := json.Unmarshal(b, &t)
	if err != nil {
		*s = NewSet[T]()
	} else {
		*s = NewSet[T](t...)
	}
	return nil
}

// GormDataType gorm common data type
func (s Set[T]) GormDataType() string {
	return "json"
}