set.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  1. package btree
  2. type Set[K ordered] struct {
  3. base Map[K, struct{}]
  4. }
  5. // Copy
  6. func (tr *Set[K]) Copy() *Set[K] {
  7. tr2 := new(Set[K])
  8. tr2.base = *tr.base.Copy()
  9. return tr2
  10. }
  11. // Insert an item
  12. func (tr *Set[K]) Insert(key K) {
  13. tr.base.Set(key, struct{}{})
  14. }
  15. func (tr *Set[K]) Scan(iter func(key K) bool) {
  16. tr.base.Scan(func(key K, value struct{}) bool {
  17. return iter(key)
  18. })
  19. }
  20. // Get a value for key
  21. func (tr *Set[K]) Contains(key K) bool {
  22. _, ok := tr.base.Get(key)
  23. return ok
  24. }
  25. // Len returns the number of items in the tree
  26. func (tr *Set[K]) Len() int {
  27. return tr.base.Len()
  28. }
  29. // Delete an item
  30. func (tr *Set[K]) Delete(key K) {
  31. tr.base.Delete(key)
  32. }
  33. // Ascend the tree within the range [pivot, last]
  34. // Pass nil for pivot to scan all item in ascending order
  35. // Return false to stop iterating
  36. func (tr *Set[K]) Ascend(pivot K, iter func(key K) bool) {
  37. tr.base.Ascend(pivot, func(key K, value struct{}) bool {
  38. return iter(key)
  39. })
  40. }
  41. func (tr *Set[K]) Reverse(iter func(key K) bool) {
  42. tr.base.Reverse(func(key K, value struct{}) bool {
  43. return iter(key)
  44. })
  45. }
  46. // Descend the tree within the range [pivot, first]
  47. // Pass nil for pivot to scan all item in descending order
  48. // Return false to stop iterating
  49. func (tr *Set[K]) Descend(pivot K, iter func(key K) bool) {
  50. tr.base.Descend(pivot, func(key K, value struct{}) bool {
  51. return iter(key)
  52. })
  53. }
  54. // Load is for bulk loading pre-sorted items
  55. func (tr *Set[K]) Load(key K) {
  56. tr.base.Load(key, struct{}{})
  57. }
  58. // Min returns the minimum item in tree.
  59. // Returns nil if the treex has no items.
  60. func (tr *Set[K]) Min() (K, bool) {
  61. key, _, ok := tr.base.Min()
  62. return key, ok
  63. }
  64. // Max returns the maximum item in tree.
  65. // Returns nil if the tree has no items.
  66. func (tr *Set[K]) Max() (K, bool) {
  67. key, _, ok := tr.base.Max()
  68. return key, ok
  69. }
  70. // PopMin removes the minimum item in tree and returns it.
  71. // Returns nil if the tree has no items.
  72. func (tr *Set[K]) PopMin() (K, bool) {
  73. key, _, ok := tr.base.PopMin()
  74. return key, ok
  75. }
  76. // PopMax removes the maximum item in tree and returns it.
  77. // Returns nil if the tree has no items.
  78. func (tr *Set[K]) PopMax() (K, bool) {
  79. key, _, ok := tr.base.PopMax()
  80. return key, ok
  81. }
  82. // GetAt returns the value at index.
  83. // Return nil if the tree is empty or the index is out of bounds.
  84. func (tr *Set[K]) GetAt(index int) (K, bool) {
  85. key, _, ok := tr.base.GetAt(index)
  86. return key, ok
  87. }
  88. // DeleteAt deletes the item at index.
  89. // Return nil if the tree is empty or the index is out of bounds.
  90. func (tr *Set[K]) DeleteAt(index int) (K, bool) {
  91. key, _, ok := tr.base.DeleteAt(index)
  92. return key, ok
  93. }
  94. // Height returns the height of the tree.
  95. // Returns zero if tree has no items.
  96. func (tr *Set[K]) Height() int {
  97. return tr.base.Height()
  98. }
  99. // SetIter represents an iterator for btree.Set
  100. type SetIter[K ordered] struct {
  101. base MapIter[K, struct{}]
  102. }
  103. // Iter returns a read-only iterator.
  104. func (tr *Set[K]) Iter() SetIter[K] {
  105. return SetIter[K]{tr.base.Iter()}
  106. }
  107. // Seek to item greater-or-equal-to key.
  108. // Returns false if there was no item found.
  109. func (iter *SetIter[K]) Seek(key K) bool {
  110. return iter.base.Seek(key)
  111. }
  112. // First moves iterator to first item in tree.
  113. // Returns false if the tree is empty.
  114. func (iter *SetIter[K]) First() bool {
  115. return iter.base.First()
  116. }
  117. // Last moves iterator to last item in tree.
  118. // Returns false if the tree is empty.
  119. func (iter *SetIter[K]) Last() bool {
  120. return iter.base.Last()
  121. }
  122. // Next moves iterator to the next item in iterator.
  123. // Returns false if the tree is empty or the iterator is at the end of
  124. // the tree.
  125. func (iter *SetIter[K]) Next() bool {
  126. return iter.base.Next()
  127. }
  128. // Prev moves iterator to the previous item in iterator.
  129. // Returns false if the tree is empty or the iterator is at the beginning of
  130. // the tree.
  131. func (iter *SetIter[K]) Prev() bool {
  132. return iter.base.Prev()
  133. }
  134. // Key returns the current iterator item key.
  135. func (iter *SetIter[K]) Key() K {
  136. return iter.base.Key()
  137. }
  138. // Keys returns all the keys in order.
  139. func (tr *Set[K]) Keys() []K {
  140. return tr.base.Keys()
  141. }