Wednesday, 05 January 2011 17:52
"""
Module for balanced binary search trees by D. Cortesi
License: Creative Commons Attribution-Noncommercial-ShareAlike
Usage: from bbst import *
defines two classes, bbstree and bbsnode.
tree = bbstree() creates a new, empty tree.
tnode = tree.lookup(k) searches tree for a node with key k.
Returns the bbsnode with key k, or None if k is not found.
tnode.key, tnode.value are usable attributes when found.
tree.forward() is a generator that returns each tree node
in ascending key order.
tree.reverse() is a generator that returns each tree node
in descending key order.
for tnode in tree.forward/reverse():
print(tnode.key, tnode.value)
tree.height() returns the height of the tree, 0 for the empty tree.
The tree's population is between 2**height and 2**(height+1),
that is, if height is 7, the tree has from 128 to 255 nodes.
The expected search length is height, the maximum is height+2.
tree.insert(k, v=None) creates a node with key=k and value=v
and inserts it into tree. If a node with key=k already
exists, its value is updated to v. Nothing is returned.
tree.instest(k, v=None) looks for a node k in the tree, and inserts
it only if it does not exist. Returns the node. Use this for
example if you are using the tree to keep a tally of items k:
tree.instest(k,0) installs node k only if it didn't exist, but
returns it, so tree.instest(k,0).value += 1 keeps a tally.
tree.delete(k) deletes the tree node with key k if it exists, and
rebalances the tree as necessary.
Any data type may be used for a key (not just integers). However, during
lookup, insert, and delete, the key k may be compared to any node's key.
Thus all keys must support ==, >, and < comparisons and be mutually
comparable by Python type rules.
The value of a node may be any Python type. To associate a value with
a key, present it as v in insert(k,v) or instest(k,v). To retrieve a
value, access tree.lookup(k).value or tree.instest(k,v).value. To change
the value of an existing key, assign to tree.lookup(k).value or to
tree.instest(k,v).value, or simply call tree.insert(k,newvalue)
The tree is balanced; adding and deleting keys in any sequence will
not produce an unbalanced tree. The algorithms are based on the
discussion by Julienne Walker found at
http://eternallyconfuzzled.com/tuts/datastructures/jsw_tut_andersson.aspx
"""
#
# A "node" is an object containing a key, a value, a level, and
# links to at most two other nodes.
#
# Walker's algorithms expect the subtrees of any leaf node to be filled
# with pointers to a nil node that acts as a sentinel. It has level=0
# and left and right subs that point to itself. Here we create that one
# node in the bbst module namespace.
#
class bbsnode_0:
def __init__(self):
self.level = 0
self.key = None
self.value = None
self.subs = [None,None]
#
NIL_NODE = bbsnode_0()
NIL_NODE.subs = [NIL_NODE,NIL_NODE]
#
# Here we create an official object of type node, with subtrees initialized
# to NIL_NODE:
#
class bbsnode:
def __init__(self,k,v=None):
self.level = 1
self.key = k
self.value = v
self.subs = [NIL_NODE,NIL_NODE]
#
# End of class bbsnode.
#
class bbstree:
#
# A "tree" is an object that contains zero or more "nodes." The tree,
# which is the "self" argument to all these member functions, is not
# itself a node, but rather contains nodes, or actually, links to the
# one node which is currently the root of the tree. That in turn links
# to sub-nodes and they to sub-sub-nodes etc.
#
# Here we initialize a new tree:
#
def __init__(self):
self.root = None # the tree is empty
self.lastinsert = None # save last-inserted node
#
# Height
#
def height(self):
if self.root: # not empty
return self.root.level
else:
return 0
#
# Look for key k in the tree; if found, return its node, else return None.
# tree.lookup() is the "public" method; it handles the special case of
# the empty tree. We use __lookup() to recurse into the tree.
#
def lookup(self, k):
if self.root: # we are not an empty tree
return self.__lookup(self.root, k)
else: # this tree has no nodes
return None # not found
#
# Recursive lookup
#
def __lookup(self, tnode, k):
if tnode.key == k:
return tnode
else:
sub = tnode.subs[k > tnode.key] # k is in left or right?
if sub != NIL_NODE: # we have a subtree that side
return self.__lookup(sub, k)
else:
return None
#
# Generator functions, forward() for stepping through the tree in order, and
# reverse() backwards. Self is a tree object, and the generator really needs to
# operate on the tree nodes, so the real generator is __scan. To scan forward,
# one goes first down the left (0) side then down the right (1) side. To scan
# backward, the reverse. Since the code is the same except for the subscripts
# of tnode.subs[], those numbers are passed in as arguments to __scan().
# The scanning function is not recursive as this would be awkward for a python
# generator (yes you can have recursive generators but a generator that calls
# itself only returns another generator object!). It emulates recursion with a
# stack. Stack size is not an issue with a balanced tree; if we stack 31
# tnodes we have a tree of 2 giga-nodes and stack size is not our main issue!
# An empty tree has self.root = None, which terminates the generator immediately.
#
def forward(self):
return self.__scan(self.root,0,1)
def reverse(self):
return self.__scan(self.root,1,0)
def __scan(self,tnode,first,second):
direction = 0
stack = [None] # when popped, ends the loop
while tnode != None:
if direction == 0:
while tnode.subs[first] != NIL_NODE:
stack.append(tnode)
tnode = tnode.subs[first]
yield tnode
if tnode.subs[second] != NIL_NODE:
tnode = tnode.subs[second]
direction = 0
else:
tnode = stack.pop()
direction = 1
#
# The following are sub-functions of __insert, used to balance the tree.
# One basic operation is the skew: "A skew removes left horizontal links by
# rotating right at the parent." In the example, a/b/c/d/e are key values
# and the digits are level values. Skew(node-with-d):
#
# d,2 b,2
# / \ / \
# b,2 e,1 --> a,1 d,2
# / \ / \
# a,1 c,1 c,1 e,1
#
# The following is Walker's non-recursive skew. Although it gets a self
# arg by virtue of being a class attribute of tree, "self" (the tree object)
# is ignored. The function operates only on the argument (a node) t.
#
def __skew(self, t):
if t.level != 0:
if t.subs[0].level == t.level:
tmp = t.subs[0]
t.subs[0] = tmp.subs[1]
tmp.subs[1] = t
t = tmp
return t
#
# The second basic operation is the split: "A split removes consecutive
# horizontal links by rotating left and increasing the level of the parent."
# Continuing the example following the above skew, split(node-with-b)
#
# b,2 d,3
# / \ / \
# a,1 d,2 --> b,2 e,2
# / \ / \
# c,1 e,2 a,1 c,1
#
# The following is Walker's non-recursive split. Again self (a tree) is ignored.
#
def __split(self, t):
if t.level != 0:
if t.level == t.subs[1].subs[1].level:
tmp = t.subs[1]
t.subs[1] = tmp.subs[0]
tmp.subs[0] = t
t = tmp
t.level += 1
return t
#
# With that, we can insert key k into the tree (or simply find it, if
# it already exists). Nothing is returned. As with tree.lookup(),
# tree.insert() is the public method and handles the special case of
# the empty tree; then __insert() is called to do the recursion.
#
def insert(self, k, v=None):
if self.root: # we are not an empty tree
self.root = self.__split( self.__skew( self.__insert(self.root, k, v) ) )
else: # empty tree getting its first node
self.root = bbsnode(k,v)
self.lastinsert = self.root
#
# Recursive insertion
#
def __insert(self, tnode, k, v):
if k == tnode.key: # key exists, update value
tnode.value = v
else: # it goes in the left or the right
side = k > tnode.key
sub = tnode.subs[ side ]
if sub != NIL_NODE: # there is a subtree that side
tnode.subs[side] = self.__split( self.__skew( self.__insert(sub, k, v) ) )
else: # no subtree on that side, make new leaf
self.lastinsert = bbsnode(k,v)
tnode.subs[side] = self.lastinsert
return tnode
#
# Because of balancing, it is impractical for insert() to return the inserted
# node: __insert() may return a different node as a result of balancing, and
# insert() doesn't know if self.lastinsert was updated or not. This makes it
# awkward to use the tree e.g. to keep a running tally of input symbols, where
# each symbol is stored as a key and the tally is the value. Here we provide
# instest() for this purpose: it can know that self.lastinsert was updated
# because of the failure of the preceding lookup.
#
def instest(self, k, v):
tnode = self.lookup(k)
if tnode == None:
self.insert(k, v)
tnode = self.lastinsert
return tnode
#
# Deletion, still following the lead of Julienne Walker although, you know,
# rewriting pretty freely. Again the public method is part of a tree and
# handles the special case of the empty tree. __delete() does the recursion.
#
def delete(self, k):
if self.root: # we are not an empty tree
self.root = self.__delete(self.root, k)
#
# Recursive deletion with nonrecursive skew and split. This is the second
# version Walker gives, described as "simple, but unconventional enough
# to be confusing." Yup. See her discussion for why it is always enough
# to do three skews and two splits. Also recall that the skew and split
# routines are conditional and do nothing but a test when not needed.
#
def __delete(self, tnode, k):
if tnode != NIL_NODE:
if k == tnode.key:
if (tnode.subs[0] != NIL_NODE) & (tnode.subs[1] != NIL_NODE) :
heir = tnode.subs[0]
while heir.subs[1] != NIL_NODE:
heir = heir.subs[1]
tnode.key = heir.key
tnode.value = heir.value
tnode.subs[0] = self.__delete(tnode.subs[0], tnode.key)
else:
tnode = tnode.subs[tnode.subs[0]==NIL_NODE]
else:
side = tnode.key < k
tnode.subs[side] = self.__delete(tnode.subs[side], k)
lv = tnode.level-1
if (tnode.subs[0].level < lv) | (tnode.subs[1].level < lv) :
tnode.level = lv
if tnode.subs[1].level > lv:
tnode.subs[1].level = lv
tnode = self.__skew(tnode)
tnode.subs[1] = self.__skew(tnode.subs[1])
tnode.subs[1].subs[1] = self.__skew(tnode.subs[1].subs[1])
tnode = self.__split(tnode)
tnode.subs[1] = self.__split(tnode.subs[1])
return tnode
Read more: http://feeds.dzone.com/~r/dzone/snippets/~3/zkXjhcxV8Nc/12811