# pylint: disable=C0302,trailing-whitespace,assigning-non-slot
# C0302: Unfortunately, we want a lot of lines in this module
# trailing-whitespace: We use trailing whitespace to keep line numbers
# consistent between the set and dict-like versions.
# assigning-non-slot: Pylint is misdetecting property assignments as
# non-slot assignments
# duplicate-code: We m4 preprocess this file into 2 or more versions, so
# of course we have "duplicate" code. But things aren't duplicated in
# the m4. Unfortunately, pylint doesn't allow us to override this
# one, so we use a regex to this-pylint in addition to this disable.
# too-many-lines: We want something self-contained, so this has a lot of
# lines
'''Red-Black tree and plain Binary Tree module'''
##Copyright (c) 2013 duncan g. smith and Dan Stromberg
##
##(This is the well-known MIT license)
##
##Permission is hereby granted, free of charge, to any person obtaining a
##copy of this software and associated documentation files (the "Software"),
##to deal in the Software without restriction, including without limitation
##the rights to use, copy, modify, merge, publish, distribute, sublicense,
##and/or sell copies of the Software, and to permit persons to whom the
##Software is furnished to do so, subject to the following conditions:
##
##The above copyright notice and this permission notice shall be included
##in all copies or substantial portions of the Software.
##
##THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
##OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
##FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
##THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
##OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
##ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
##OTHER DEALINGS IN THE SOFTWARE.
# This code was originally by Duncan G. Smith, but has been modified a
# bit by Dan Stromberg.
changequote(`/*', `*/')
from __future__ import division
import sys
import math
import random
import itertools
if hasattr(itertools, 'izip_longest'):
MY_ZIP_LONGEST = getattr(itertools, 'izip_longest')
else:
MY_ZIP_LONGEST = getattr(itertools, 'zip_longest')
try:
next
except NameError:
def next(iterator):
# pylint: disable=redefined-builtin
'''A version of next() for python's that don't have it'''
return iterator.next()
#class TreeError(Exception):
# '''An exception to raise if the tree gets in a bad state - unused'''
# pass
def center(string, field_use_width, field_avail_width):
'''Center a string within a given field width'''
field_use = (string + '_' * (field_use_width - 1))[:field_use_width - 1]
pad_width = (field_avail_width - /*len*/(field_use)) / 2.0
result = ' ' * int(pad_width) + field_use + ' ' * int(math.ceil(pad_width))
return result
class BinaryTree(object):
# pylint: disable=too-many-public-methods
"""
A basic binary tree class. Arbitrary data can be
associated with a tree. A tree that is root has
parent equal to None; a tree that is leaf has left
and right children that are empty trees (sentinels).
An empty tree has left and right children equal to None.
A tree can have children (or a parent) with equal
data.
"""
__slots__ = ('key', ifdef(/*m4keyvalue*/,/*'value', */)'left', 'right', 'parent')
def __init__(self, key=None, ifdef(/*m4keyvalue*/,/*value=None, */)left=None, right=None, parent=None):
"""
Initialises instance.
@type key: arbitrary type
@param key: data
ifdef(/*m4keyvalue*/,/*@type value: arbitrary type*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*@param value: data*/,/*# placeholder*/)
@type left: L{BinaryTree} or C{None}
@param left: left child
@type right: L{BinaryTree} or C{None}
@param right: right child
@type parent: L{BinaryTree} or C{None}
@param parent: parent
"""
self.key = key
ifdef(/*m4keyvalue*/,/*self.value = value*/,/*# placeholder*/)
self.left = left
self.right = right
self.parent = parent
def __repr__(self):
"""
Returns a string representation of I{self}, containing
the key in I{self} and left and right children. If I{self}
is a sentinel the string "sentinel" is returned.
@rtype: C{str}
@return: string representation
"""
if self:
return '/*%s*/ifdef(/*m4keyvalue*/,/* %s*/)' % (self.key, ifdef(/*m4keyvalue*/,/*self.value*/))
else:
return "sent"
def _depth_and_field_width(self):
'''Compute the depth of the tree and the maximum width within the nodes - for pretty printing'''
class maxer(object):
'''Class facilitates computing the max depth of the treap (tree) and max width of the nodes'''
def __init__(self, maximum=-1):
self.depth_max = maximum
self.field_width_max = -1
def feed(self, node, key, ifdef(/*m4keyvalue*/,/*value, */)depth, from_left):
'''Check our maximums so far against the current node - updating as needed'''
# pylint: disable=R0913
# R0913: We need a bunch of arguments
dummy = key
ifdef(/*m4keyvalue*/,/*dummy = value*/,/*# placeholder*/)
dummy = from_left
if depth > self.depth_max:
self.depth_max = depth
repr_node = repr(node)
len_node = /*len*/(repr_node)
if len_node > self.field_width_max:
self.field_width_max = len_node
def result(self):
'''Return the maximums we've computed'''
return (self.depth_max + 1, self.field_width_max)
max_obj = maxer()
self.detailed_inorder_traversal(max_obj.feed)
return max_obj.result()
def to_dot(self, initial=True, file_=sys.stdout, visited=None):
"""Generate a dot file describing this tree"""
if visited is None:
visited = set()
if initial:
file_.write('digraph G {\n')
this_node = '%s %s' % (id(self), repr(self))
if this_node in visited:
return
visited.add(this_node)
if self.left is not None:
file_.write(' "%s" -> "%s %s" [ label="left" ]\n' % (this_node, id(self.left), repr(self.left)))
self.left.to_dot(initial=False, file_=file_, visited=visited)
if self.right is not None:
file_.write(' "%s" -> "%s %s" [ label="right" ]\n' % (this_node, id(self.right), repr(self.right)))
self.right.to_dot(initial=False, file_=file_, visited=visited)
if initial:
file_.write('}\n')
def __str__(self):
'''Format a tree as a string'''
class Desc(object):
# pylint: disable=R0903
# R0903: We don't need a lot of public methods
'''Build a pretty-print string during a recursive traversal'''
def __init__(self, pretree):
self.tree = pretree
def update(self, node, key, ifdef(/*m4keyvalue*/,/*value, */)depth, from_left):
'''Add a node to the tree'''
# pylint: disable=R0913
# R0913: We need a bunch of arguments
dummy = key
ifdef(/*m4keyvalue*/,/*dummy = value*/,/*# placeholder*/)
self.tree[depth][from_left] = repr(node)
pretree = []
depth, field_use_width = self._depth_and_field_width()
field_use_width += 1
for level in range(depth):
string = '_' * (field_use_width - 1)
pretree.append([string] * 2 ** level)
desc = Desc(pretree)
self.detailed_inorder_traversal(desc.update)
result = []
widest = 2 ** (depth - 1) * field_use_width
for level in range(depth):
two_level = 2 ** level
field_avail_width = widest / two_level
string = ''.join([center(x, field_use_width, field_avail_width) for x in desc.tree[level]])
# this really isn't useful for more than dozen values or so
result.append(('%2d ' % level) + string)
return '\n'.join(result)
def __nonzero__(self):
"""
A sentinel node evaluates to False.
@rtype: C{bool}
@return: True if neither left nor right are None,
otherwise False
"""
return not (self.left is None and self.right is None)
__bool__ = __nonzero__
@property
def is_root(self):
"""
Returns True if I{self} is root.
@rtype: C{bool}
@return: True if I{self} is root, otherwise False
"""
return self.parent is None
@property
def sibling(self):
"""
Returns the sibling of I{self}.
@rtype: L{BinaryTree} or C{None}
@return: sibling
"""
# sibling may be a sentinel node
if self.is_root:
return None
elif self is self.parent.left:
return self.parent.right
else:
return self.parent.left
@property
def size(self):
"""
Return the number of key items / non-sentinel
nodes in the subtree rooted at I{self}.
@rtype: C{int}
@return: number of elements
"""
if self:
return self.left.size + self.right.size + 1
return 0
@property
def num_levels(self):
"""
Returns the number of levels in the subtree
rooted at I{self}.
@rtype: C{int}
@return: number of levels
"""
if self:
return max(self.left.num_levels, self.right.num_levels) + 1
return 0
@property
def height(self):
"""
Returns the height of the subtree rooted at
I{self}. This is the number of levels - 1.
@rtype: C{int}
@return: number of levels
"""
return self.num_levels - 1
def copy(self, parent=None):
"""
Returns a shallow copy of the subtree rooted
at I{self} (does not copy data).
@type parent: L{BinaryTree} or C{None}
@param parent: parent
@rtype: L{BinaryTree}
@return: a copy
"""
if self:
copy = self.__class__(key=self.key, ifdef(/*m4keyvalue*/,/*value=self.value, */)parent=parent)
copy.left = self.left.copy(copy)
copy.right = self.right.copy(copy)
else:
copy = self.__class__(parent=parent)
return copy
def del_node(self):
"""
Deletes a node from the subtree rooted at I{self},
but does not delete the subtree rooted at node. To
delete a node on the basis of its data use the
I{remove} method.
"""
if self.left:
if self.right:
# method chosen randomly in order to
# prevent tree becoming too unbalanced
if random.random() < 0.5:
node = self.left.maximum
self.key = node.key
ifdef(/*m4keyvalue*/,/*self.value = node.value*/,/*# placeholder*/)
node.del_node()
else:
node = self.right.minimum
self.key = node.key
ifdef(/*m4keyvalue*/,/*self.value = node.value*/,/*# placeholder*/)
node.del_node()
else:
node = self.left
self.key, self.left, self.right = node.key, node.left, node.right
ifdef(/*m4keyvalue*/,/*self.value = node.value*/,/*# placeholder*/)
self.left.parent = self
self.right.parent = self
elif self.right:
node = self.right
self.key, self.left, self.right = node.key, node.left, node.right
ifdef(/*m4keyvalue*/,/*self.value = node.value*/,/*# placeholder*/)
self.left.parent = self
self.right.parent = self
else:
# make node a sentinel node
self.key = self.left = self.right = None
ifdef(/*m4keyvalue*/,/*self.value = None*/,/*# placeholder*/)
def add(self, /*key*/ifdef(/*m4keyvalue*/,/*, value*/)):
"""
Adds a node containing I{key} to the subtree
rooted at I{self}, returning the added node.
@type key: arbitrary type
@param key: key
ifdef(/*m4keyvalue*/,/*@type value: arbitrary type*/,/*placeholder*/)
ifdef(/*m4keyvalue*/,/*@param value: data*/,/*placeholder*/)
@rtype: L{BinaryTree}
@return: (replaced flag, added node)
"""
node = self.find(key)
if not node:
node.key = key
ifdef(/*m4keyvalue*/,/*node.value = value*/,/*# placeholder*/)
node.left, node.right = self.__class__(parent=node), self.__class__(parent=node)
return (False, node)
ifdef(/*m4set*/,/*elif node.key == key:*/,/*# placeholder*/)
ifdef(/*m4set*/,/*node.key = key*/,/*# placeholder*/)
ifelse(m4set/**/m4keyvalue,/*11*/,/*node.value = value*/,/*# placeholder*/)
ifdef(/*m4set*/,/*return (True, node)*/,/*# placeholder*/)
else:
if random.random() < 0.5:
return BinaryTree.add(node.left, /*key=key*/ifdef(/*m4keyvalue*/,/*, value=value*/))
else:
return BinaryTree.add(node.right, /*key=key*/ifdef(/*m4keyvalue*/,/*, value=value*/))
def remove(self, key):
"""
Removes a node containing I{key} from the
subtree rooted at I{self}, raising an error
if the subtree does not contain I{key}.
@type key: arbitrary type
@param key: data
@raise ValueError: if key is not contained in the
subtree rooted at I{self}
"""
node = self.find(key)
if node:
node.del_node()
else:
raise ValueError('tree.remove(x): x not in tree')
def remove_min(self):
"""
Removes minimum node from subtree rooted at I{self},
raising an exception if the subtree does not have any
nodes.
@raise ValueError: if key is not contained in the
subtree rooted at I{self}
"""
node = self.minimum
if node:
# If this is a dict-like-object, return a tuple of the node.key and
# node.value. If this is a set-like-object, return the just key
ifdef(/*m4keyvalue*/,/*result = (node.key, node.value)*/,/*result = node.key*/)
node.del_node()
return result
else:
raise ValueError('tree.remove(x): x not in tree')
def remove_max(self):
"""
Removes maximum node from subtree rooted at I{self},
raising an exception if the subtree does not have any
nodes.
@raise ValueError: if key is not contained in the
subtree rooted at I{self}
"""
node = self.maximum
if node:
# If this is a dict-like-object, return a tuple of the node.key and
# node.value. If this is a set-like-object, return the just key
ifdef(/*m4keyvalue*/,/*result = (node.key, node.value)*/,/*result = node.key*/)
node.del_node()
return result
else:
raise ValueError('tree.remove(x): x not in tree')
def discard(self, key):
"""
Discards a node containing I{key} from the subtree
rooted at I{self} (without raising an error if the key
is not present).
@type key: arbitrary type
@param key: data
"""
try:
self.remove(key)
except ValueError:
pass
def count(self, key):
"""
Returns the number of occurrences of I{key} in
the subtree rooted at I{self}.
@type key: arbitrary type
@param key: data
@rtype: C{int}
@return: data count
"""
if self:
node = self.find(key)
if node:
return node.left.count(key) + node.right.count(key) + 1
return 0
def __add__(self, iterable):
"""
Returns a tree with elements from the subtree
rooted at I{self} and I{iterable}. I{iterable}
may be any iterable, not necessarily another
tree.
@type iterable: any iterable type
@param iterable: an iterable object
@rtype: L{BinaryTree}
@return: tree containing the elements in I{self}
and iterable
"""
items = list(self) + list(iterable)
random.shuffle(items) # try to avoid unbalanced tree
tree = self.__class__()
for item in items:
tree.add(item)
return tree
def __iadd__(self, iterable):
"""
Adds the elements in I{iterable} to the
subtree rooted at I{self}.
@type iterable: any iterable type
@param iterable: an iterable object
"""
items = list(iterable)
random.shuffle(items) # try to avoid unbalanced tree
for item in items:
self.add(item)
def find(self, key):
"""
Finds and returns a node containing I{key} in the
subtree rooted at I{self}. If I{key} is not in
the tree, then a sentinel in the location where
I{key} can be inserted is returned.
@type key: arbitrary type
@param key: data
@rtype: L{BinaryTree}
@return: node containing key, or sentinel
node
"""
if self:
if self.key == key:
return self
elif key < self.key:
return self.left.find(key)
else:
return self.right.find(key)
return self
@property
def predecessor(self):
"""
Returns the predecessor of I{self}, or None if
self has no predecessor.
@rtype: L{BinaryTree} or None
@return: predecessor of I{self}
"""
if self.left:
return self.left.maximum
else:
if not self.parent:
return None
if self is self.parent.right:
return self.parent
else:
current, parent = self, self.parent
while parent and parent.left is current:
current, parent = parent, parent.parent
return parent
@property
def successor(self):
"""
Returns the successor of I{self}, or None if
self has no successor.
@rtype: L{BinaryTree} or None
@return: successor of I{self}
"""
if self.right:
return self.right.minimum
else:
if not self.parent:
return None
if self is self.parent.left:
return self.parent
else:
current, parent = self, self.parent
while parent and parent.right is current:
current, parent = parent, parent.parent
return parent
def in_order(self):
"""
Returns a generator of nodes in the subtree
rooted at I{self} in sorted order of node keys.
@rtype: C{generator}
@return: generator of sorted nodes
"""
if self:
if self.left:
for node in self.left.in_order():
yield node
yield self
if self.right:
for node in self.right.in_order():
yield node
def detailed_inorder_traversal(self, visit, depth=0, from_left=0):
'''Do an inorder traversal - with lots of parameters'''
if self.left:
self.left.detailed_inorder_traversal(visit, depth + 1, from_left * 2)
visit(self, self.key, ifdef(/*m4keyvalue*/,/*self.value, */)depth, from_left)
if self.right:
self.right.detailed_inorder_traversal(visit, depth + 1, from_left * 2 + 1)
def __iter__(self):
"""
Returns a generator of the node keys in
the subtree rooted at I{self} in sorted
order.
@rtype: C{generator}
@return: generator of sorted node values
"""
return (n.key for n in self.in_order())
def in_reverse(self):
"""
Returns a generator of nodes in the subtree
rooted at I{self} in reverse sorted order
of node keys.
@rtype: C{generator}
@return: generator of reverse sorted nodes
"""
if self:
for node in self.right.in_reverse():
yield node
yield self
for node in self.left.in_reverse():
yield node
def pre_order(self):
"""
Returns a generator of the nodes in the
subtree rooted at I{self} in preorder.
@rtype: C{generator}
@return: generator of nodes in preorder
"""
if self:
yield self
if self.left:
for node in self.left.pre_order():
yield node
if self.right:
for node in self.right.pre_order():
yield node
def post_order(self):
"""
Returns a generator of the nodes in the
subtree rooted at I{self} in postorder.
@rtype: C{generator}
@return: generator of nodes in postorder
"""
if self:
if self.left:
for node in self.left.post_order():
yield node
if self.right:
for node in self.right.post_order():
yield node
yield self
@property
def minimum(self):
"""
Returns a node for which the node key is minimum
in the subtree rooted at I{self}. Always returns
the minimum-valued node with no left child.
@rtype: L{BinaryTree}
@return: node with minimum key value
"""
node = self
while node.left:
node = node.left
return node
@property
def maximum(self):
"""
Returns a node for which the node key is maximum
in the subtree rooted at I{self}. Always returns
the maximum-valued node with no right child.
@rtype: L{BinaryTree}
@return: node with maximum key value
"""
node = self
while node.right:
node = node.right
return node
def find_min(self):
'''
Return the miniumum key in the tree
'''
node = self.minimum
return node.key
def find_max(self):
'''
Return the maxiumum key in the tree
'''
node = self.maximum
return node.key
def is_similar(self, other):
"""
Two binary trees are similar if they have the
same structure.
@rtype: C{bool}
@return: True if the subtrees rooted at I{self}
and I{other} are similar, otherwise False
"""
if not self and not other:
return True
if self and other:
return self.left.is_similar(other.left) and self.right.is_similar(other.right)
return False
def is_equivalent(self, other):
"""
Two binary trees are equivalent if they are similar
and corresponding nodes contain the same keys.
@rtype: C{bool}
@return: True if the subtrees rooted at I{self}
and I{other} are equivalent, otherwise False
"""
if not self and not other:
return True
if self and other:
return (
self.key == other.key and
ifdef(/*m4keyvalue*/,/*self.value == other.value and*/,/*# placeholder*/)
self.left.is_equivalent(other.left) and
self.right.is_equivalent(other.right)
)
return False
# def cmp(self, other):
# """
# Compares I{self} with I{other} lexocographically.
#
# @type other: L{BinaryTree}
# @param other: tree being compared to I{self}
# @rtype: C{int}
# @return: 0 if the subtrees rooted at I{self}
# and I{other} contain equal values,
# -1 if the ordered values in I{self} are lexicographically
# less than the ordered values in I{other}, and
# 1 if the ordered values in I{self} are lexicographically
# greater than the ordered values in I{other}
# """
# other_items = iter(other)
# for self_item in self:
# try:
# other_item = next(other_items)
# except StopIteration:
# return 1
# else:
# if self_item < other_item:
# return -1
# elif self_item > other_item:
# return 1
# try:
# next(other_items)
# except StopIteration:
# return 0
def __cmp__(self, other):
"""
Returns C{0} if the subtrees rooted at
I{self} and I{other} contain equal keys.
I{other} need not be similar to I{self}.
@type other: L{BinaryTree}
@param other: tree being compared to I{self}
@rtype: C{int}
@return: 0 if the subtrees rooted at I{self}
and I{other} do contain equal values,
-1 if self < other.
1 if self > other.
"""
# FIXME: This probably should compare values too.
for self_node, other_node in MY_ZIP_LONGEST(self.in_order(), other.in_order()):
if self_node is None:
return -1
elif other_node is None:
return 1
elif self_node.key < other_node.key:
return -1
elif self_node.key > other_node.key:
return 1
return 0
cmp = __cmp__
def __eq__(self, other):
"""
Returns C{True} if the subtrees rooted at
I{self} and I{other} contain equal keys.
I{other} need not be similar to I{self}.
@type other: L{BinaryTree}
@param other: tree being compared to I{self}
@rtype: C{bool}
@return: True if the subtrees rooted at I{self}
and I{other} contain equal values, otherwise
False
"""
return self.cmp(other) == 0
def __neq__(self, other):
"""
Returns C{True} if the subtrees rooted at
I{self} and I{other} do not contain equal keys.
I{other} need not be similar to I{self}.
@type other: L{BinaryTree}
@param other: tree being compared to I{self}
@rtype: C{bool}
@return: True if the subtrees rooted at I{self}
and I{other} do not contain equal values,
otherwise False
"""
return self.cmp(other) != 0
def __lt__(self, other):
"""
Returns C{True} if the subtrees rooted at
I{self} and I{other} do not contain equal keys.
I{other} need not be similar to I{self}.
@type other: L{BinaryTree}
@param other: tree being compared to I{self}
@rtype: C{bool}
@return: True if the subtrees rooted at I{self}
and I{other} do not contain equal values,
otherwise False
"""
return self.cmp(other) < 0
def __gt__(self, other):
return other.cmp(self) < 0
def __le__(self, other):
return self.cmp(other) <= 0
def __ge__(self, other):
return self.cmp(other) >= 0
def __contains__(self, key):
"""
Returns true if I{key} is stored in any node in
the subtree rooted at I{self}.
@type key: arbitrary type
@param key: data
@rtype: C{bool}
@return: True if I{key} in subtree rooted at
I{self}, otherwise False
"""
return bool(self.find(key))
class RedBlackTree(BinaryTree):
# pylint: disable=attribute-defined-outside-init,maybe-no-member,too-many-public-methods
"""
A binary tree with an extra attribute, is_red. Colour (red
or black) is used to maintain a reasonably balanced tree.
Changes to the tree structure (rotations) are carried out
in such a way that the name bound to the initial empty
tree will always refer to the root node. The data stored
in a given node may change, as well as its connectivity.
e.g.
>>> import trees
>>> tree = trees.RedBlackTree()
>>> tree.add(2)
2 False None None
>>> tree.key # tree is bound to root node
2
>>> tree.add(-1)
-1 True None None
>>> tree.add(1) # tree rotations are performed at this point
1 False -1 2
>>> tree.key # the root node has changed (i.e. its key has been changed)
1
>>>
"""
__slots__ = BinaryTree.__slots__ + ('is_red', )
def __init__(self, key=None, ifdef(/*m4keyvalue*/,/*value=None, */)left=None, right=None, parent=None):
"""
Initialises instance.
@type key: arbitrary type
@param key: data
@type left: L{RedBlackTree} or C{None}
@param left: left child
@type right: L{RedBlackTree} or C{None}
@param right: right child
@type parent: L{RedBlackTree} or C{None}
@param parent: parent
"""
super(RedBlackTree, self).__init__(key, ifdef(/*m4keyvalue*/,/*value, */)left, right, parent)
self.is_red = False # default for sentinel nodes
@property
def is_black(self):
"""
Returns True if node colour is black.
@rtype: C{bool}
@return: True if node colour is black, otherwise False
"""
return not self.is_red
@is_black.setter
def is_black(self, val):
"""
Sets node colour.
"""
self.is_red = not val
@property
def black_height(self):
"""
@rtype: C{int}
@return: black height of subtree rooted at I{self}
"""
if self.left:
if self.left.is_black:
return self.left.black_height + 1
else:
return self.left.black_height
elif self.right:
if self.right.is_black:
return self.right.black_height + 1
else:
return self.right.black_height
else:
return 0
def __repr__(self):
"""
Returns a string representation of I{self}, containing
the key in I{self}, the colour, and the key in left and
right children. If I{self} is a sentinel the string
"sentinel" is returned.
@rtype: C{str}
@return: string representation
"""
if self:
if self.is_red:
color_string = 'red'
else:
color_string = 'blk'
return '%s ifdef(/*m4keyvalue*/,/*%s */)%s' % (self.key, ifdef(/*m4keyvalue*/,/*self.value, */)color_string)
else:
return "sent"
def copy(self, parent=None):
"""
Returns a shallow copy of the subtree rooted
at I{self} (does not copy key).
@type parent: L{RedBlackTree} or C{None}
@param parent: parent
@rtype: L{RedBlackTree}
@return: a copy
"""
copy = super(RedBlackTree, self).copy(parent)
copy.is_red = self.is_red
return copy
def add(self, /*key*/ifdef(/*m4keyvalue*/,/*, value*/)):
"""
Adds a node containing I{key} to the subtree
rooted at I{self}, returning the node which
contains I{key}. Note: the returned node might
be an existing node due to key swaps on rotations.
@type key: arbitrary type
@param key: data
@rtype: L{RedBlackTree}
@return: node containing I{key}
"""
(replaced, node) = super(RedBlackTree, self).add(key=/*key*/ifdef(/*m4keyvalue*/,/*, value=value*/))
if not replaced:
node.is_red = True
node = node.fix_insert()
return node
def fix_insert(self):
"""
Carries out tree rotations, key swaps and recolourings
to ensure that the red / black properties of the tree
are preserved on node insertion. The returned node contains
the same value as I{self} (but might not be the inserted node
due to key swaps / rotations).
@rtype: L{RedBlackTree}
@return: node containing key originally
contained in I{self}
"""
# case 1
if self.parent is None:
self.is_black = True
return self
# case 2
if self.parent.is_black:
return self
# case 3
sib = self.parent.sibling
if sib.is_red:
self.parent.is_black = True
sib.is_black = True
par = self.parent.parent
par.is_red = True
par.fix_insert()
return self
# cases 4 & 5
par = self.parent.parent
if not (self.parent.left is self) == (par.left is self.parent):
self.rotate()
self.parent.rotate()
return par
else:
self.parent.rotate()
return self
def rotate(self):
"""
Rotates I{self} with I{self.parent}. Left / right rotations
are performed as dictated by the parent-child relationship.
The colours of I{self} and I{self.parent} are implicity swapped
(as their values are swapped during the rotation).
"""
# data swapped and connections updated
# colours are (implicitly) swapped
if self is self.parent.left:
self.key, self.parent.key = self.parent.key, self.key
ifdef(/*m4keyvalue*/,/*self.value, self.parent.value = self.parent.value, self.value*/,/*# placeholder*/)
self.parent.left = self.left
self.parent.left.parent = self.parent
self.left = self.right
self.right = self.parent.right
self.left.parent = self.right.parent = self
self.parent.right = self
else:
self.key, self.parent.key = self.parent.key, self.key
ifdef(/*m4keyvalue*/,/*self.value, self.parent.value = self.parent.value, self.value*/,/*# placeholder*/)
self.parent.right = self.right
self.parent.right.parent = self.parent
self.right = self.left
self.left = self.parent.left
self.left.parent = self.right.parent = self
self.parent.left = self
def del_node(self):
"""
Deletes a node from the subtree rooted at I{self},
but does not delete the subtree rooted at node. To
delete a node on the basis of its key use the
I{remove} method.
"""
# ensure that node being deleted has at most one child
if self.left:
if self.right:
# method chosen randomly in order to
# prevent tree becoming too unbalanced
if random.random() < 0.5:
node = self.left.maximum
else:
node = self.right.minimum
self.key, node.key = node.key, self.key
ifdef(/*m4keyvalue*/,/*self.value, node.value = node.value, self.value*/,/*# placeholder*/)
node.del_node() # node has at most one child
else:
node = self.left
self.key, self.left, self.right = node.key, node.left, node.right
ifdef(/*m4keyvalue*/,/*self.value = node.value*/,/*# placeholder*/)
self.left.parent = self # this might result in sentinel node's parent attribute being set
self.right.parent = self
self.is_black = True
elif self.right:
node = self.right
self.key, self.left, self.right = node.key, node.left, node.right
ifdef(/*m4keyvalue*/,/*self.value = node.value*/,/*# placeholder*/)
self.left.parent = self
self.right.parent = self
self.is_black = True
else:
# node being deleted is leaf
self.key = self.left = self.right = None
ifdef(/*m4keyvalue*/,/*self.value = None*/,/*# placeholder*/)
if self.is_red:# or self.sibling: # just make it a sentinel node
self.is_black = True
else:
self.is_black = True
# self has extra blackness
self.fix_del()
def fix_del(self):
"""
Maintains red / black property on node
deletion.
"""
# case 1
if self.parent is None:
return
sib = self.sibling
# case 2
if sib.is_red:
sib.rotate()
sib = self.sibling
# case 3
if (
self.parent.is_black and
sib.is_black and
sib.left.is_black and
sib.right.is_black
):
sib.is_red = True
self.parent.fix_del()
return
# case 4
if (
self.parent.is_red and
sib.is_black and
sib.left.is_black and
sib.right.is_black
):
sib.is_red = True
self.parent.is_black = True
return
# case 5
if sib.is_black:
if self is self.parent.left and sib.right.is_black and sib.left.is_red:
sib.left.rotate()
elif self is self.parent.right and sib.left.is_black and sib.right.is_red:
sib.right.rotate()
sib = self.sibling
# case 6
sib.rotate()
self.parent.sibling.is_black = True
def check(self):
"""
Checks whether the subtree rooted at I{self} is a valid
red black tree U{http://en.wikipedia.org/wiki/Red_black_tree}.
@rtype: C{bool}
@return: True if the tree is valid, otherwise False
"""
leaf_nodes = []
for node in self.in_order():
if not node.left and not node.right:
leaf_nodes.append(node)
if node.is_red:
if not (node.left.is_black and node.right.is_black):
return False
black_height = 0
if leaf_nodes:
node = leaf_nodes[0]
while node.parent:
if node.is_black:
black_height += 1
node = node.parent
for node in leaf_nodes[1:]:
height = 0
while node.parent:
if node.is_black:
height += 1
node = node.parent
if not height == black_height:
return False
return True
ifdef(/*m4keyvalue*/,/*def __getitem__(self, key):*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*'''Get an item by its key'''*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*node = self.find(key)*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*if node:*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*return node.value*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*else:*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*raise KeyError*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*def __setitem__(self, key, value):*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*'''Add an item using a dictionary-like interface'''*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*self.add(key, value)*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*def __delitem__(self, key):*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*'''Delete an item using a dictionary-like interface'''*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*node = self.find(key)*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*if node:*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*node.del_node()*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*else:*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*raise KeyError*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*def iteritems(self):*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*'''Iterate over the items in the "dictionary"'''*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*return ((node.key, node.value) for node in self.in_order())*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*items = iteritems*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*def itervalues(self):*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*'''Iterate over the values in the "dictionary"'''*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*return (node.value for node in self.in_order())*/,/*# placeholder*/)
ifdef(/*m4keyvalue*/,/*values = itervalues*/,/*# placeholder*/)
def __len__(self):
'''Return the number of elements in the container'''
return self.size
def _slow_len(self):
'''Compute the length of the tree using a slow-but-accurate method'''
count = 0
for element in self:
dummy = element
count += 1
return count