treap.go

// Copyright 2011 Numerotron Inc.
// Use of this source code is governed by an MIT-style license
// that can be found in the LICENSE file.
//
// Developed at www.stathat.com by Patrick Crosby
// Contact us on twitter with any questions:  twitter.com/stat_hat

// The treap package provides a balanced binary tree datastructure, expected
// to have logarithmic height.
package treap

import (
        "rand"
        "time"
)

func init() {
        rand.Seed(time.Nanoseconds())
}

// A Tree is the data structure that stores everything
type Tree struct {
        less    LessFunc
        overlap OverlapFunc
        count   int
        root    *Node
}

// LessFunc returns true if a < b
type LessFunc func(a, b interface{}) bool

// OverlapFunc return true if a overlaps b.  Optional.  Only used by overlap trees.
type OverlapFunc func(a, b interface{}) bool

// Key can be anything.  It will use LessFunc to compare keys.
type Key interface{}

// Item can be anything.
type Item interface{}

// A Node in the Tree.
type Node struct {
        key      Key
        item     Item
        priority int
        left     *Node
        right    *Node
}

func newNode(key Key, item Item, priority int) *Node {
        result := new(Node)
        result.key = key
        result.item = item
        result.priority = priority
        return result
}

// To create a Tree, you need to supply a LessFunc that can compare the
// keys in the Node.
func NewTree(lessfn LessFunc) *Tree {
        t := new(Tree)
        t.less = lessfn
        return t
}

// To create a tree that also lets you iterate by key overlap, supply a LessFunc
// and an OverlapFunc
func NewOverlapTree(lessfn LessFunc, overfn OverlapFunc) *Tree {
        t := new(Tree)
        t.less = lessfn
        t.overlap = overfn
        return t
}

// Remove everything from the tree.
func (t *Tree) Reset() {
        t.root = nil
        t.count = 0
}

// The length of the tree (number of nodes).
func (t *Tree) Len() int {
        return t.count
}

// Get an Item in the tree.
func (t *Tree) Get(key Key) Item {
        return t.get(t.root, key)
}

func (t *Tree) get(node *Node, key Key) Item {
        if node == nil {
                return nil
        }
        if t.less(key, node.key) {
                return t.get(node.left, key)
        }
        if t.less(node.key, key) {
                return t.get(node.right, key)
        }
        return node.item
}

// Returns true if there is an item in the tree with this key.
func (t *Tree) Exists(key Key) bool {
        return t.exists(t.root, key)
}

func (t *Tree) exists(node *Node, key Key) bool {
        if node == nil {
                return false
        }
        if t.less(key, node.key) {
                return t.exists(node.left, key)
        }
        if t.less(node.key, key) {
                return t.exists(node.right, key)
        }
        return true
}

// Insert an item into the tree.
func (t *Tree) Insert(key Key, item Item) {
        priority := rand.Int()
        t.root = t.insert(t.root, key, item, priority)
}

func (t *Tree) insert(node *Node, key Key, item Item, priority int) *Node {
        if node == nil {
                t.count++
                return newNode(key, item, priority)
        }
        if t.less(key, node.key) {
                node.left = t.insert(node.left, key, item, priority)
                if node.left.priority < node.priority {
                        return t.leftRotate(node)
                }
                return node
        }
        if t.less(node.key, key) {
                node.right = t.insert(node.right, key, item, priority)
                if node.right.priority < node.priority {
                        return t.rightRotate(node)
                }
                return node
        }

        // equal: replace the value
        node.item = item
        return node
}

func (t *Tree) leftRotate(node *Node) *Node {
        result := node.left
        x := result.right
        result.right = node
        node.left = x
        return result
}

func (t *Tree) rightRotate(node *Node) *Node {
        result := node.right
        x := result.left
        result.left = node
        node.right = x
        return result
}

// Split the tree by creating a tree with a node of priority -1 so it will be the root
func (t *Tree) Split(key Key) (*Node, *Node) {
        inserted := t.insert(t.root, key, nil, -1)
        return inserted.left, inserted.right
}

// Merge two trees together by supplying the root node of each tree.
func (t *Tree) Merge(left, right *Node) *Node {
        if left == nil {
                return right
        }
        if right == nil {
                return left
        }
        if left.priority < right.priority {
                result := left
                x := left.right
                result.right = t.Merge(x, right)
                return result
        }

        result := right
        x := right.left
        result.left = t.Merge(x, left)
        return result
}

// Delete the item from the tree that has this key.
func (t *Tree) Delete(key Key) {
        if t.Exists(key) == false {
                return
        }
        t.root = t.delete(t.root, key)
}

func (t *Tree) delete(node *Node, key Key) *Node {
        if node == nil {
                panic("key not found")
        }

        if t.less(key, node.key) {
                result := node
                x := node.left
                result.left = t.delete(x, key)
                return result
        }
        if t.less(node.key, key) {
                result := node
                x := node.right
                result.right = t.delete(x, key)
                return result
        }
        t.count--
        return t.Merge(node.left, node.right)
}

// Returns the height (depth) of the tree.
func (t *Tree) Height(key Key) int {
        return t.height(t.root, key)
}

func (t *Tree) height(node *Node, key Key) int {
        if node == nil {
                return 0
        }
        if t.less(key, node.key) {
                depth := t.height(node.left, key)
                return depth + 1
        }
        if t.less(node.key, key) {
                depth := t.height(node.right, key)
                return depth + 1
        }
        return 0
}

// Returns a channel of Items whose keys are in ascending order.
func (t *Tree) IterAscend() <-chan Item {
        c := make(chan Item)
        go func() {
                iterateInOrder(t.root, c)
                close(c)
        }()
        return c
}

func iterateInOrder(h *Node, c chan<- Item) {
        if h == nil {
                return
        }
        iterateInOrder(h.left, c)
        c <- h.item
        iterateInOrder(h.right, c)
}

// Returns a channel of keys in ascending order.
func (t *Tree) IterKeysAscend() <-chan Key {
        c := make(chan Key)
        go func() {
                iterateKeysInOrder(t.root, c)
                close(c)
        }()
        return c
}

func iterateKeysInOrder(h *Node, c chan<- Key) {
        if h == nil {
                return
        }
        iterateKeysInOrder(h.left, c)
        c <- h.key
        iterateKeysInOrder(h.right, c)
}

// Returns a channel of items that overlap key.
func (t *Tree) IterateOverlap(key Key) <-chan Item {
        c := make(chan Item)
        go func() {
                if t.overlap != nil {
                        t.iterateOverlap(t.root, key, c)
                }
                close(c)
        }()
        return c
}

func (t *Tree) iterateOverlap(h *Node, key Key, c chan<- Item) {
        if h == nil {
                return
        }
        lessThanLower := t.overlap(h.key, key)
        greaterThanUpper := t.overlap(key, h.key)

        if !lessThanLower {
                t.iterateOverlap(h.left, key, c)
        }
        if !lessThanLower && !greaterThanUpper {
                c <- h.item
        }
        if !greaterThanUpper {
                t.iterateOverlap(h.right, key, c)
        }
}

// Returns the minimum item in the tree.
func (t *Tree) Min() Item {
        return min(t.root)
}

func min(h *Node) Item {
        if h == nil {
                return nil
        }
        if h.left == nil {
                return h.item
        }
        return min(h.left)
}

// Returns the maximum item in the tree.
func (t *Tree) Max() Item {
        return max(t.root)
}

func max(h *Node) Item {
        if h == nil {
                return nil
        }
        if h.right == nil {
                return h.item
        }
        return max(h.right)
}