You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

188 lines
4.0 KiB

  1. package gtreap
  2. type Treap struct {
  3. compare Compare
  4. root *node
  5. }
  6. // Compare returns an integer comparing the two items
  7. // lexicographically. The result will be 0 if a==b, -1 if a < b, and
  8. // +1 if a > b.
  9. type Compare func(a, b interface{}) int
  10. // Item can be anything.
  11. type Item interface{}
  12. type node struct {
  13. item Item
  14. priority int
  15. left *node
  16. right *node
  17. }
  18. func NewTreap(c Compare) *Treap {
  19. return &Treap{compare: c, root: nil}
  20. }
  21. func (t *Treap) Min() Item {
  22. n := t.root
  23. if n == nil {
  24. return nil
  25. }
  26. for n.left != nil {
  27. n = n.left
  28. }
  29. return n.item
  30. }
  31. func (t *Treap) Max() Item {
  32. n := t.root
  33. if n == nil {
  34. return nil
  35. }
  36. for n.right != nil {
  37. n = n.right
  38. }
  39. return n.item
  40. }
  41. func (t *Treap) Get(target Item) Item {
  42. n := t.root
  43. for n != nil {
  44. c := t.compare(target, n.item)
  45. if c < 0 {
  46. n = n.left
  47. } else if c > 0 {
  48. n = n.right
  49. } else {
  50. return n.item
  51. }
  52. }
  53. return nil
  54. }
  55. // Note: only the priority of the first insert of an item is used.
  56. // Priorities from future updates on already existing items are
  57. // ignored. To change the priority for an item, you need to do a
  58. // Delete then an Upsert.
  59. func (t *Treap) Upsert(item Item, itemPriority int) *Treap {
  60. r := t.union(t.root, &node{item: item, priority: itemPriority})
  61. return &Treap{compare: t.compare, root: r}
  62. }
  63. func (t *Treap) union(this *node, that *node) *node {
  64. if this == nil {
  65. return that
  66. }
  67. if that == nil {
  68. return this
  69. }
  70. if this.priority > that.priority {
  71. left, middle, right := t.split(that, this.item)
  72. if middle == nil {
  73. return &node{
  74. item: this.item,
  75. priority: this.priority,
  76. left: t.union(this.left, left),
  77. right: t.union(this.right, right),
  78. }
  79. }
  80. return &node{
  81. item: middle.item,
  82. priority: this.priority,
  83. left: t.union(this.left, left),
  84. right: t.union(this.right, right),
  85. }
  86. }
  87. // We don't use middle because the "that" has precendence.
  88. left, _, right := t.split(this, that.item)
  89. return &node{
  90. item: that.item,
  91. priority: that.priority,
  92. left: t.union(left, that.left),
  93. right: t.union(right, that.right),
  94. }
  95. }
  96. // Splits a treap into two treaps based on a split item "s".
  97. // The result tuple-3 means (left, X, right), where X is either...
  98. // nil - meaning the item s was not in the original treap.
  99. // non-nil - returning the node that had item s.
  100. // The tuple-3's left result treap has items < s,
  101. // and the tuple-3's right result treap has items > s.
  102. func (t *Treap) split(n *node, s Item) (*node, *node, *node) {
  103. if n == nil {
  104. return nil, nil, nil
  105. }
  106. c := t.compare(s, n.item)
  107. if c == 0 {
  108. return n.left, n, n.right
  109. }
  110. if c < 0 {
  111. left, middle, right := t.split(n.left, s)
  112. return left, middle, &node{
  113. item: n.item,
  114. priority: n.priority,
  115. left: right,
  116. right: n.right,
  117. }
  118. }
  119. left, middle, right := t.split(n.right, s)
  120. return &node{
  121. item: n.item,
  122. priority: n.priority,
  123. left: n.left,
  124. right: left,
  125. }, middle, right
  126. }
  127. func (t *Treap) Delete(target Item) *Treap {
  128. left, _, right := t.split(t.root, target)
  129. return &Treap{compare: t.compare, root: t.join(left, right)}
  130. }
  131. // All the items from this are < items from that.
  132. func (t *Treap) join(this *node, that *node) *node {
  133. if this == nil {
  134. return that
  135. }
  136. if that == nil {
  137. return this
  138. }
  139. if this.priority > that.priority {
  140. return &node{
  141. item: this.item,
  142. priority: this.priority,
  143. left: this.left,
  144. right: t.join(this.right, that),
  145. }
  146. }
  147. return &node{
  148. item: that.item,
  149. priority: that.priority,
  150. left: t.join(this, that.left),
  151. right: that.right,
  152. }
  153. }
  154. type ItemVisitor func(i Item) bool
  155. // Visit items greater-than-or-equal to the pivot.
  156. func (t *Treap) VisitAscend(pivot Item, visitor ItemVisitor) {
  157. t.visitAscend(t.root, pivot, visitor)
  158. }
  159. func (t *Treap) visitAscend(n *node, pivot Item, visitor ItemVisitor) bool {
  160. if n == nil {
  161. return true
  162. }
  163. if t.compare(pivot, n.item) <= 0 {
  164. if !t.visitAscend(n.left, pivot, visitor) {
  165. return false
  166. }
  167. if !visitor(n.item) {
  168. return false
  169. }
  170. }
  171. return t.visitAscend(n.right, pivot, visitor)
  172. }