blob: f758ffe444b427616e5b32668e3361f649d4a937 [file] [log] [blame]
package gtreap
type Treap struct {
compare Compare
root *node
}
// Compare returns an integer comparing the two items
// lexicographically. The result will be 0 if a==b, -1 if a < b, and
// +1 if a > b.
type Compare func(a, b interface{}) int
// Item can be anything.
type Item interface{}
type node struct {
item Item
priority int
left *node
right *node
}
func NewTreap(c Compare) *Treap {
return &Treap{compare: c, root: nil}
}
func (t *Treap) Min() Item {
n := t.root
if n == nil {
return nil
}
for n.left != nil {
n = n.left
}
return n.item
}
func (t *Treap) Max() Item {
n := t.root
if n == nil {
return nil
}
for n.right != nil {
n = n.right
}
return n.item
}
func (t *Treap) Get(target Item) Item {
n := t.root
for n != nil {
c := t.compare(target, n.item)
if c < 0 {
n = n.left
} else if c > 0 {
n = n.right
} else {
return n.item
}
}
return nil
}
// Note: only the priority of the first insert of an item is used.
// Priorities from future updates on already existing items are
// ignored. To change the priority for an item, you need to do a
// Delete then an Upsert.
func (t *Treap) Upsert(item Item, itemPriority int) *Treap {
r := t.union(t.root, &node{item: item, priority: itemPriority})
return &Treap{compare: t.compare, root: r}
}
func (t *Treap) union(this *node, that *node) *node {
if this == nil {
return that
}
if that == nil {
return this
}
if this.priority > that.priority {
left, middle, right := t.split(that, this.item)
if middle == nil {
return &node{
item: this.item,
priority: this.priority,
left: t.union(this.left, left),
right: t.union(this.right, right),
}
}
return &node{
item: middle.item,
priority: this.priority,
left: t.union(this.left, left),
right: t.union(this.right, right),
}
}
// We don't use middle because the "that" has precendence.
left, _, right := t.split(this, that.item)
return &node{
item: that.item,
priority: that.priority,
left: t.union(left, that.left),
right: t.union(right, that.right),
}
}
// Splits a treap into two treaps based on a split item "s".
// The result tuple-3 means (left, X, right), where X is either...
// nil - meaning the item s was not in the original treap.
// non-nil - returning the node that had item s.
// The tuple-3's left result treap has items < s,
// and the tuple-3's right result treap has items > s.
func (t *Treap) split(n *node, s Item) (*node, *node, *node) {
if n == nil {
return nil, nil, nil
}
c := t.compare(s, n.item)
if c == 0 {
return n.left, n, n.right
}
if c < 0 {
left, middle, right := t.split(n.left, s)
return left, middle, &node{
item: n.item,
priority: n.priority,
left: right,
right: n.right,
}
}
left, middle, right := t.split(n.right, s)
return &node{
item: n.item,
priority: n.priority,
left: n.left,
right: left,
}, middle, right
}
func (t *Treap) Delete(target Item) *Treap {
left, _, right := t.split(t.root, target)
return &Treap{compare: t.compare, root: t.join(left, right)}
}
// All the items from this are < items from that.
func (t *Treap) join(this *node, that *node) *node {
if this == nil {
return that
}
if that == nil {
return this
}
if this.priority > that.priority {
return &node{
item: this.item,
priority: this.priority,
left: this.left,
right: t.join(this.right, that),
}
}
return &node{
item: that.item,
priority: that.priority,
left: t.join(this, that.left),
right: that.right,
}
}
type ItemVisitor func(i Item) bool
// Visit items greater-than-or-equal to the pivot.
func (t *Treap) VisitAscend(pivot Item, visitor ItemVisitor) {
t.visitAscend(t.root, pivot, visitor)
}
func (t *Treap) visitAscend(n *node, pivot Item, visitor ItemVisitor) bool {
if n == nil {
return true
}
if t.compare(pivot, n.item) <= 0 {
if !t.visitAscend(n.left, pivot, visitor) {
return false
}
if !visitor(n.item) {
return false
}
}
return t.visitAscend(n.right, pivot, visitor)
}