github.com/go-board/x-go@v0.1.2-0.20220610024734-db1323f6cb15/xcontainer/set/string_btree_set.go (about)

     1  package set
     2  
     3  import (
     4  	"github.com/google/btree"
     5  )
     6  
     7  type StringItem string
     8  
     9  func (s StringItem) Less(than btree.Item) bool {
    10  	return s < than.(StringItem)
    11  }
    12  
    13  type StringBtreeSet struct {
    14  	tree *btree.BTree
    15  }
    16  
    17  func NewStringBtreeSet(items ...string) *StringBtreeSet {
    18  	tree := btree.New(64)
    19  	for _, item := range items {
    20  		tree.ReplaceOrInsert(StringItem(item))
    21  	}
    22  	set := &StringBtreeSet{tree: tree}
    23  	return set
    24  }
    25  
    26  func (s *StringBtreeSet) Add(items ...string) {
    27  	for _, item := range items {
    28  		s.tree.ReplaceOrInsert(StringItem(item))
    29  	}
    30  }
    31  
    32  func (s *StringBtreeSet) Remove(items ...string) {
    33  	for _, item := range items {
    34  		s.tree.Delete(StringItem(item))
    35  	}
    36  }
    37  
    38  func (s *StringBtreeSet) Values() []string {
    39  	values := make([]string, 0, s.tree.Len())
    40  	s.tree.Ascend(func(i btree.Item) bool {
    41  		values = append(values, string(i.(StringItem)))
    42  		return true
    43  	})
    44  	return values
    45  }
    46  
    47  func (s *StringBtreeSet) RangeAll(fn func(item string)) {
    48  	s.tree.Ascend(func(i btree.Item) bool {
    49  		fn(string(i.(StringItem)))
    50  		return true
    51  	})
    52  }
    53  
    54  func (s *StringBtreeSet) Contains(item string) bool {
    55  	return s.tree.Has(StringItem(item))
    56  }
    57  
    58  func (s *StringBtreeSet) ContainAll(items ...string) bool {
    59  	for _, item := range items {
    60  		if !s.tree.Has(StringItem(item)) {
    61  			return false
    62  		}
    63  	}
    64  	return true
    65  }
    66  
    67  func (s *StringBtreeSet) ContainAny(items ...string) bool {
    68  	for _, item := range items {
    69  		if s.tree.Has(StringItem(item)) {
    70  			return true
    71  		}
    72  	}
    73  	return false
    74  }
    75  
    76  func (s *StringBtreeSet) Size() int {
    77  	return s.tree.Len()
    78  }
    79  
    80  func (s *StringBtreeSet) Union(o *StringBtreeSet) *StringBtreeSet {
    81  	n := NewStringBtreeSet(s.Values()...)
    82  	o.RangeAll(func(item string) {
    83  		n.Add(item)
    84  	})
    85  	return n
    86  }
    87  
    88  func (s *StringBtreeSet) InterSection(o *StringBtreeSet) *StringBtreeSet {
    89  	n := NewStringBtreeSet()
    90  	s.RangeAll(func(item string) {
    91  		if o.ContainAll(item) {
    92  			n.Add(item)
    93  		}
    94  	})
    95  	return n
    96  }
    97  
    98  func (s *StringBtreeSet) Difference(o *StringBtreeSet) *StringBtreeSet {
    99  	n := NewStringBtreeSet(s.Values()...)
   100  	n.Remove(o.Values()...)
   101  	return n
   102  }
   103  
   104  func (s *StringBtreeSet) IsSuperSet(o *StringBtreeSet) bool {
   105  	return s.ContainAll(o.Values()...)
   106  }
   107  
   108  func (s *StringBtreeSet) IsSubSet(o *StringBtreeSet) bool {
   109  	return o.ContainAll(s.Values()...)
   110  }
   111  
   112  func (s *StringBtreeSet) Range(begin, end string, fn func(key string) bool) {
   113  	s.tree.AscendGreaterOrEqual(StringItem(begin), func(i btree.Item) bool {
   114  		item := i.(StringItem)
   115  		return string(item) < end && fn(string(item))
   116  	})
   117  }