返回顶部

收藏

a balanced binary search tree data structure

更多

来源:https://github.com/stathat/treap

This Go package provides a balanced binary search tree data structure, expected to have logarithmic height.

This implementation borrows a lot of ideas from GoLLRB by Petar Maymounkov.

// Copyright 2011 Numrotron 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 (
    "math/rand"
    "time"
)

func init() {
    rand.Seed(time.Now().UnixNano())
}

// 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)
}
//该片段来自于http://outofmemory.cn

标签:golang,基础

收藏

0人收藏

支持

0

反对

0

»更多 您可能感兴趣的代码
  1. 2012-11-05 21:59:42java获得随机数代码 by 怪兽狂殴奥特曼
  2. 2017-03-31 15:12:31Catch several exceptions by jack.chen
  3. 2012-11-06 01:19:02获得hashtable的集合视图 by DuLinRain
  4. 2012-11-08 17:54:01java类型转换 by 朱凯迪
  5. 2012-11-19 21:54:37java使用Vector例子 by Liner
  6. 2013-01-31 21:25:11java使用DecimalFormat格式化数字 by 迟浩东
  7. 2013-02-07 20:43:32java使用Character类 by 司马
  8. 2013-07-25 12:48:05Linq基本语法 by 迟浩东
  9. 2014-02-19 17:51:50go-lisp , 不足20行给你一个lisp交互解释器 by river0830
  10. 2014-03-09 10:14:12图像与缓存 by 屁屁果
  11. 2014-03-18 15:35:04日志的按天分类压缩保存 by owensh