use generic version of btree (#3404)

This commit is contained in:
Denis Mishin 2022-06-06 14:31:05 -04:00 committed by GitHub
parent 63673363c5
commit f7b6ed0ad4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -2,51 +2,47 @@ package sets
import "github.com/google/btree"
type stringItem string
func (item stringItem) Less(than btree.Item) bool {
return item < than.(stringItem)
}
func lessFn(a, b string) bool { return a < b }
// A SortedString is a set of strings with sorted iteration.
type SortedString struct {
b *btree.BTree
b *btree.BTreeG[string]
}
// NewSortedString creates a new sorted string set.
func NewSortedString() *SortedString {
return &SortedString{
b: btree.New(8),
b: btree.NewG[string](8, lessFn),
}
}
// Add adds a string to the set.
func (s *SortedString) Add(elements ...string) {
for _, element := range elements {
s.b.ReplaceOrInsert(stringItem(element))
s.b.ReplaceOrInsert(element)
}
}
// Clear clears the set.
func (s *SortedString) Clear() {
s.b = btree.New(8)
s.b = btree.NewG[string](8, lessFn)
}
// Delete deletes an element from the set.
func (s *SortedString) Delete(element string) {
s.b.Delete(stringItem(element))
s.b.Delete(element)
}
// ForEach iterates over the set in ascending order.
func (s *SortedString) ForEach(callback func(element string) bool) {
s.b.Ascend(func(i btree.Item) bool {
return callback(string(i.(stringItem)))
s.b.Ascend(func(item string) bool {
return callback(item)
})
}
// Has returns true if the elment is in the set.
func (s *SortedString) Has(element string) bool {
return s.b.Has(stringItem(element))
return s.b.Has(element)
}
// Size returns the size of the set.
@ -57,8 +53,8 @@ func (s *SortedString) Size() int {
// ToSlice returns a slice of all the elements in the set.
func (s *SortedString) ToSlice() []string {
arr := make([]string, 0, s.Size())
s.b.Ascend(func(i btree.Item) bool {
arr = append(arr, string(i.(stringItem)))
s.b.Ascend(func(item string) bool {
arr = append(arr, item)
return true
})
return arr