Mercurial > public > think_complexity

annotate redblacktree.py @ 18:92e2879e2e33

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Rework the red-black tree based on Julienne Walker's tutorial. Insertion is implemented now. Deletion will come next.
author Brian Neal <bgneal@gmail.com>
date Wed, 26 Dec 2012 19:59:17 -0600
parents 977628018b4b
children 3c74185c5047
rev   line source
bgneal@18 1 """
bgneal@15 2 Copyright (C) 2012 Brian G. Neal.
bgneal@15 3
bgneal@15 4 Permission is hereby granted, free of charge, to any person obtaining a copy of
bgneal@15 5 this software and associated documentation files (the "Software"), to deal in
bgneal@15 6 the Software without restriction, including without limitation the rights to
bgneal@15 7 use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
bgneal@15 8 the Software, and to permit persons to whom the Software is furnished to do so,
bgneal@15 9 subject to the following conditions:
bgneal@15 10
bgneal@15 11 The above copyright notice and this permission notice shall be included in all
bgneal@15 12 copies or substantial portions of the Software.
bgneal@15 13
bgneal@15 14 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
bgneal@15 15 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
bgneal@15 16 FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
bgneal@15 17 COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
bgneal@15 18 IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
bgneal@15 19 CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
bgneal@15 20
bgneal@18 21 ----
bgneal@18 22
bgneal@18 23 A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think
bgneal@18 24 Complexity_ book.
bgneal@18 25
bgneal@18 26 http://greenteapress.com/complexity
bgneal@18 27
bgneal@18 28 Red black trees are described on Wikipedia:
bgneal@18 29 http://en.wikipedia.org/wiki/Red-black_tree.
bgneal@18 30
bgneal@18 31 This is basically a Python implementation of Julienne Walker's Red Black Trees
bgneal@18 32 tutorial found at:
bgneal@18 33 http://www.eternallyconfuzzled.com/tuts/datastructures/jsw_tut_rbtree.aspx
bgneal@18 34 We implement Julienne's top-down insertion and deletion algorithms here.
bgneal@18 35
bgneal@18 36 Some ideas were also borrowed from code by Darren Hart at
bgneal@18 37 http://dvhart.com/darren/files/rbtree.py
bgneal@18 38
bgneal@15 39 """
bgneal@15 40
bgneal@16 41 BLACK, RED = range(2)
bgneal@18 42 LEFT, RIGHT = range(2)
bgneal@18 43
bgneal@18 44
bgneal@18 45 class TreeError(Exception):
bgneal@18 46 """Base exception class for red-black tree errors."""
bgneal@18 47 pass
bgneal@18 48
bgneal@15 49
bgneal@15 50 class Node(object):
bgneal@15 51 """A node class for red-black trees.
bgneal@15 52
bgneal@18 53 * A node has a color, either RED or BLACK.
bgneal@18 54 * A node has a key and an optional value.
bgneal@15 55
bgneal@18 56 * The key is used to order the red-black tree by calling the "<"
bgneal@18 57 operator when comparing two nodes.
bgneal@18 58 * The value is useful for using the red-black tree to implement a map
bgneal@18 59 datastructure.
bgneal@18 60
bgneal@18 61 * Nodes have exactly 2 link fields which we represent as a list of
bgneal@18 62 2 elements to represent the left and right children of the node. This list
bgneal@18 63 representation was borrowed from Julienne Walker as it simplifies some of
bgneal@18 64 the code. Element 0 is the LEFT child and element 1 is the RIGHT child.
bgneal@15 65
bgneal@15 66 """
bgneal@15 67
bgneal@18 68 def __init__(self, key, value=None, color=RED):
bgneal@15 69 self.key = key
bgneal@15 70 self.value = value
bgneal@15 71 self.color = color
bgneal@18 72 self.link = [None, None]
bgneal@17 73
bgneal@17 74 def __str__(self):
bgneal@17 75 c = 'B' if self.color == BLACK else 'R'
bgneal@17 76 if self.value:
bgneal@17 77 return '({}: {} => {})'.format(c, self.key, self.value)
bgneal@17 78 else:
bgneal@17 79 return '({}: {})'.format(c, self.key)
bgneal@17 80
bgneal@17 81
bgneal@18 82 def is_red(n):
bgneal@18 83 """Return True if the given Node n is RED and False otherwise.
bgneal@18 84
bgneal@18 85 If the node is None, then it is considered to be BLACK, and False is
bgneal@18 86 returned.
bgneal@18 87
bgneal@18 88 """
bgneal@18 89 return n is not None and n.color == RED
bgneal@18 90
bgneal@18 91
bgneal@17 92 class Tree(object):
bgneal@17 93 """A red-black Tree class.
bgneal@17 94
bgneal@17 95 A red-black tree is a binary search tree with the following properties:
bgneal@17 96
bgneal@17 97 1. A node is either red or black.
bgneal@17 98 2. The root is black.
bgneal@17 99 3. All leaves are black.
bgneal@17 100 4. Both children of every red node are black.
bgneal@17 101 5. Every simple path from a given node to any descendant leaf contains
bgneal@17 102 the same number of black nodes.
bgneal@17 103
bgneal@17 104 These rules ensure that the path from the root to the furthest leaf is no
bgneal@17 105 more than twice as long as the path from the root to the nearest leaf. Thus
bgneal@17 106 the tree is roughly height-balanced.
bgneal@17 107
bgneal@17 108 """
bgneal@17 109
bgneal@17 110 def __init__(self):
bgneal@17 111 self.root = None
bgneal@17 112
bgneal@17 113 def __iter__(self):
bgneal@17 114 return self._inorder(self.root)
bgneal@17 115
bgneal@17 116 def _inorder(self, node):
bgneal@17 117 """A generator to perform an inorder traversal of the nodes in the
bgneal@17 118 tree starting at the given node.
bgneal@17 119
bgneal@17 120 """
bgneal@18 121 if node.link[0]:
bgneal@18 122 for n in self._inorder(node.link[0]):
bgneal@17 123 yield n
bgneal@17 124
bgneal@17 125 yield node
bgneal@17 126
bgneal@18 127 if node.link[1]:
bgneal@18 128 for n in self._inorder(node.link[1]):
bgneal@17 129 yield n
bgneal@17 130
bgneal@18 131 def _single_rotate(self, root, d):
bgneal@18 132 """Perform a single rotation about the node 'root' in the given
bgneal@18 133 direction 'd' (LEFT or RIGHT).
bgneal@18 134
bgneal@18 135 The old root is set to RED and the new root is set to BLACK.
bgneal@18 136
bgneal@18 137 Returns the new root.
bgneal@18 138
bgneal@18 139 """
bgneal@18 140 nd = not d
bgneal@18 141 save = root.link[nd]
bgneal@18 142
bgneal@18 143 root.link[nd] = save.link[d]
bgneal@18 144 save.link[d] = root
bgneal@18 145
bgneal@18 146 root.color = RED
bgneal@18 147 save.color = BLACK
bgneal@18 148
bgneal@18 149 return save
bgneal@18 150
bgneal@18 151 def _double_rotate(self, root, d):
bgneal@18 152 """Perform two single rotations about the node root in the direction d.
bgneal@18 153
bgneal@18 154 The new root is returned.
bgneal@18 155
bgneal@18 156 """
bgneal@18 157 nd = not d
bgneal@18 158 root.link[nd] = self._single_rotate(root.link[nd], nd)
bgneal@18 159 return self._single_rotate(root, d)
bgneal@18 160
bgneal@18 161 def validate(self, root=None):
bgneal@18 162 """Checks to see if the red-black tree validates at the given node, i.e.
bgneal@18 163 all red-black tree rules are valid. If root is None, the root of the
bgneal@18 164 tree is used as the starting point.
bgneal@18 165
bgneal@18 166 If any rules are violated, a TreeError is raised.
bgneal@18 167
bgneal@18 168 Returns the black height of the tree rooted at root.
bgneal@18 169
bgneal@18 170 """
bgneal@18 171 if root is None:
bgneal@18 172 root = self.root
bgneal@18 173
bgneal@18 174 return self._validate(root)
bgneal@18 175
bgneal@18 176 def _validate(self, root):
bgneal@18 177 """Internal implementation of the validate() method."""
bgneal@18 178
bgneal@18 179 if root is None:
bgneal@18 180 return 1
bgneal@18 181
bgneal@18 182 ln = root.link[0]
bgneal@18 183 rn = root.link[1]
bgneal@18 184
bgneal@18 185 # Check for consecutive red links
bgneal@18 186
bgneal@18 187 if is_red(root) and (is_red(ln) or is_red(rn)):
bgneal@18 188 raise TreeError('red violation')
bgneal@18 189
bgneal@18 190 lh = self._validate(ln)
bgneal@18 191 rh = self._validate(rn)
bgneal@18 192
bgneal@18 193 # Check for invalid binary search tree
bgneal@18 194
bgneal@18 195 if (ln and ln.key >= root.key) or (rn and rn.key <= root.key):
bgneal@18 196 raise TreeError('binary tree violation')
bgneal@18 197
bgneal@18 198 # Check for black height mismatch
bgneal@18 199 if lh != rh:
bgneal@18 200 raise TreeError('black violation')
bgneal@18 201
bgneal@18 202 # Only count black links
bgneal@18 203
bgneal@18 204 return lh if is_red(root) else lh + 1
bgneal@18 205
bgneal@17 206 def insert(self, key, value=None):
bgneal@18 207 """Insert the (key, value) pair into the tree."""
bgneal@17 208
bgneal@18 209 # Check for the empty tree case:
bgneal@17 210 if self.root is None:
bgneal@18 211 self.root = Node(key=key, value=value, color=BLACK)
bgneal@17 212 return
bgneal@17 213
bgneal@18 214 # False/dummy tree root
bgneal@18 215 head = Node(key=None, value=None, color=BLACK)
bgneal@18 216 d = LEFT
bgneal@18 217 last = LEFT
bgneal@17 218
bgneal@18 219 t = head
bgneal@18 220 g = p = None
bgneal@18 221 t.link[1] = self.root
bgneal@18 222 q = self.root
bgneal@17 223
bgneal@18 224 # Search down the tree
bgneal@18 225 while True:
bgneal@18 226 if q is None:
bgneal@18 227 # Insert new node at the bottom
bgneal@18 228 p.link[d] = q = Node(key=key, value=value, color=RED)
bgneal@18 229 elif is_red(q.link[0]) and is_red(q.link[1]):
bgneal@18 230 # Color flip
bgneal@18 231 q.color = RED
bgneal@18 232 q.link[0].color = BLACK
bgneal@18 233 q.link[1].color = BLACK
bgneal@17 234
bgneal@18 235 # Fix red violation
bgneal@18 236 if is_red(q) and is_red(p):
bgneal@18 237 d2 = t.link[1] is g
bgneal@17 238
bgneal@18 239 if q is p.link[last]:
bgneal@18 240 t.link[d2] = self._single_rotate(g, not last)
bgneal@18 241 else:
bgneal@18 242 t.link[d2] = self._double_rotate(g, not last)
bgneal@17 243
bgneal@18 244 # Stop if found
bgneal@18 245 if q.key == key:
bgneal@17 246 break
bgneal@17 247
bgneal@18 248 last = d
bgneal@18 249 d = q.key < key
bgneal@17 250
bgneal@18 251 # Update helpers
bgneal@18 252 if g is not None:
bgneal@18 253 t = g;
bgneal@18 254 g = p
bgneal@18 255 p = q
bgneal@18 256 q = q.link[d]
bgneal@17 257
bgneal@18 258 # Update root
bgneal@18 259 self.root = head.link[1]
bgneal@18 260 self.root.color = BLACK
bgneal@17 261
bgneal@17 262
bgneal@17 263 if __name__ == '__main__':
bgneal@17 264 import random
bgneal@17 265
bgneal@18 266 for n in range(30):
bgneal@18 267 tree = Tree()
bgneal@18 268 tree.validate()
bgneal@17 269
bgneal@18 270 for i in range(25):
bgneal@18 271 val = random.randint(0, 100)
bgneal@18 272 tree.insert(val)
bgneal@17 273
bgneal@18 274 for n in tree:
bgneal@18 275 print n.key,
bgneal@18 276 print
bgneal@18 277
bgneal@18 278 tree.validate()