Mercurial > public > think_complexity

annotate redblacktree.py @ 20:0326803882ad

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Finally completing Ch. 3, exercise 5: write a TreeMap that uses a red-black tree.
author Brian Neal <bgneal@gmail.com>
date Thu, 27 Dec 2012 15:30:49 -0600
parents 3c74185c5047
children
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@19 74 def free(self):
bgneal@19 75 """Call this function when removing a node from a tree.
bgneal@19 76
bgneal@19 77 It updates the links to encourage garbage collection.
bgneal@19 78
bgneal@19 79 """
bgneal@19 80 self.link[0] = self.link[1] = None
bgneal@19 81
bgneal@17 82 def __str__(self):
bgneal@17 83 c = 'B' if self.color == BLACK else 'R'
bgneal@17 84 if self.value:
bgneal@17 85 return '({}: {} => {})'.format(c, self.key, self.value)
bgneal@17 86 else:
bgneal@17 87 return '({}: {})'.format(c, self.key)
bgneal@17 88
bgneal@17 89
bgneal@18 90 def is_red(n):
bgneal@18 91 """Return True if the given Node n is RED and False otherwise.
bgneal@18 92
bgneal@18 93 If the node is None, then it is considered to be BLACK, and False is
bgneal@18 94 returned.
bgneal@18 95
bgneal@18 96 """
bgneal@18 97 return n is not None and n.color == RED
bgneal@18 98
bgneal@18 99
bgneal@17 100 class Tree(object):
bgneal@17 101 """A red-black Tree class.
bgneal@17 102
bgneal@17 103 A red-black tree is a binary search tree with the following properties:
bgneal@17 104
bgneal@17 105 1. A node is either red or black.
bgneal@17 106 2. The root is black.
bgneal@17 107 3. All leaves are black.
bgneal@17 108 4. Both children of every red node are black.
bgneal@17 109 5. Every simple path from a given node to any descendant leaf contains
bgneal@17 110 the same number of black nodes.
bgneal@17 111
bgneal@17 112 These rules ensure that the path from the root to the furthest leaf is no
bgneal@17 113 more than twice as long as the path from the root to the nearest leaf. Thus
bgneal@17 114 the tree is roughly height-balanced.
bgneal@17 115
bgneal@17 116 """
bgneal@17 117
bgneal@17 118 def __init__(self):
bgneal@17 119 self.root = None
bgneal@19 120 self._size = 0
bgneal@19 121
bgneal@19 122 def __len__(self):
bgneal@19 123 """To support the len() function by returning the number of elements in
bgneal@19 124 the tree.
bgneal@19 125
bgneal@19 126 """
bgneal@19 127 return self._size
bgneal@17 128
bgneal@17 129 def __iter__(self):
bgneal@19 130 """Return an iterator to perform an inorder traversal of the tree."""
bgneal@17 131 return self._inorder(self.root)
bgneal@17 132
bgneal@17 133 def _inorder(self, node):
bgneal@17 134 """A generator to perform an inorder traversal of the nodes in the
bgneal@17 135 tree starting at the given node.
bgneal@17 136
bgneal@17 137 """
bgneal@19 138 if node:
bgneal@19 139 if node.link[0]:
bgneal@19 140 for n in self._inorder(node.link[0]):
bgneal@19 141 yield n
bgneal@17 142
bgneal@19 143 yield node
bgneal@17 144
bgneal@19 145 if node.link[1]:
bgneal@19 146 for n in self._inorder(node.link[1]):
bgneal@19 147 yield n
bgneal@17 148
bgneal@18 149 def _single_rotate(self, root, d):
bgneal@18 150 """Perform a single rotation about the node 'root' in the given
bgneal@18 151 direction 'd' (LEFT or RIGHT).
bgneal@18 152
bgneal@18 153 The old root is set to RED and the new root is set to BLACK.
bgneal@18 154
bgneal@18 155 Returns the new root.
bgneal@18 156
bgneal@18 157 """
bgneal@18 158 nd = not d
bgneal@18 159 save = root.link[nd]
bgneal@18 160
bgneal@18 161 root.link[nd] = save.link[d]
bgneal@18 162 save.link[d] = root
bgneal@18 163
bgneal@18 164 root.color = RED
bgneal@18 165 save.color = BLACK
bgneal@18 166
bgneal@18 167 return save
bgneal@18 168
bgneal@18 169 def _double_rotate(self, root, d):
bgneal@18 170 """Perform two single rotations about the node root in the direction d.
bgneal@18 171
bgneal@18 172 The new root is returned.
bgneal@18 173
bgneal@18 174 """
bgneal@18 175 nd = not d
bgneal@18 176 root.link[nd] = self._single_rotate(root.link[nd], nd)
bgneal@18 177 return self._single_rotate(root, d)
bgneal@18 178
bgneal@18 179 def validate(self, root=None):
bgneal@18 180 """Checks to see if the red-black tree validates at the given node, i.e.
bgneal@18 181 all red-black tree rules are valid. If root is None, the root of the
bgneal@18 182 tree is used as the starting point.
bgneal@18 183
bgneal@18 184 If any rules are violated, a TreeError is raised.
bgneal@18 185
bgneal@18 186 Returns the black height of the tree rooted at root.
bgneal@18 187
bgneal@18 188 """
bgneal@18 189 if root is None:
bgneal@18 190 root = self.root
bgneal@18 191
bgneal@18 192 return self._validate(root)
bgneal@18 193
bgneal@18 194 def _validate(self, root):
bgneal@18 195 """Internal implementation of the validate() method."""
bgneal@18 196
bgneal@18 197 if root is None:
bgneal@18 198 return 1
bgneal@18 199
bgneal@18 200 ln = root.link[0]
bgneal@18 201 rn = root.link[1]
bgneal@18 202
bgneal@18 203 # Check for consecutive red links
bgneal@18 204
bgneal@18 205 if is_red(root) and (is_red(ln) or is_red(rn)):
bgneal@18 206 raise TreeError('red violation')
bgneal@18 207
bgneal@18 208 lh = self._validate(ln)
bgneal@18 209 rh = self._validate(rn)
bgneal@18 210
bgneal@18 211 # Check for invalid binary search tree
bgneal@18 212
bgneal@18 213 if (ln and ln.key >= root.key) or (rn and rn.key <= root.key):
bgneal@18 214 raise TreeError('binary tree violation')
bgneal@18 215
bgneal@18 216 # Check for black height mismatch
bgneal@18 217 if lh != rh:
bgneal@18 218 raise TreeError('black violation')
bgneal@18 219
bgneal@18 220 # Only count black links
bgneal@18 221
bgneal@18 222 return lh if is_red(root) else lh + 1
bgneal@18 223
bgneal@17 224 def insert(self, key, value=None):
bgneal@19 225 """Insert the (key, value) pair into the tree. Duplicate keys are not
bgneal@19 226 allowed in the tree.
bgneal@19 227
bgneal@19 228 Returns a tuple of the form (node, flag) where node is either the newly
bgneal@19 229 inserted tree node or an already existing node that has the same key.
bgneal@19 230 The flag member is a Boolean that will be True if the node was inserted
bgneal@19 231 and False if a node with the given key already exists in the tree.
bgneal@19 232
bgneal@19 233 """
bgneal@17 234
bgneal@18 235 # Check for the empty tree case:
bgneal@17 236 if self.root is None:
bgneal@18 237 self.root = Node(key=key, value=value, color=BLACK)
bgneal@19 238 self._size = 1
bgneal@19 239 return self.root, True
bgneal@17 240
bgneal@18 241 # False/dummy tree root
bgneal@18 242 head = Node(key=None, value=None, color=BLACK)
bgneal@19 243 d = last = LEFT # direction variables
bgneal@17 244
bgneal@19 245 # Set up helpers
bgneal@18 246 t = head
bgneal@18 247 g = p = None
bgneal@19 248 q = t.link[1] = self.root
bgneal@19 249
bgneal@19 250 # Return values
bgneal@19 251 target, insert_flag = (None, False)
bgneal@17 252
bgneal@18 253 # Search down the tree
bgneal@18 254 while True:
bgneal@18 255 if q is None:
bgneal@18 256 # Insert new node at the bottom
bgneal@18 257 p.link[d] = q = Node(key=key, value=value, color=RED)
bgneal@19 258 self._size += 1
bgneal@19 259 insert_flag = True
bgneal@18 260 elif is_red(q.link[0]) and is_red(q.link[1]):
bgneal@18 261 # Color flip
bgneal@18 262 q.color = RED
bgneal@18 263 q.link[0].color = BLACK
bgneal@18 264 q.link[1].color = BLACK
bgneal@17 265
bgneal@18 266 # Fix red violation
bgneal@18 267 if is_red(q) and is_red(p):
bgneal@18 268 d2 = t.link[1] is g
bgneal@17 269
bgneal@18 270 if q is p.link[last]:
bgneal@18 271 t.link[d2] = self._single_rotate(g, not last)
bgneal@18 272 else:
bgneal@18 273 t.link[d2] = self._double_rotate(g, not last)
bgneal@17 274
bgneal@18 275 # Stop if found
bgneal@18 276 if q.key == key:
bgneal@19 277 target = q
bgneal@17 278 break
bgneal@17 279
bgneal@18 280 last = d
bgneal@18 281 d = q.key < key
bgneal@17 282
bgneal@18 283 # Update helpers
bgneal@18 284 if g is not None:
bgneal@18 285 t = g;
bgneal@18 286 g = p
bgneal@18 287 p = q
bgneal@18 288 q = q.link[d]
bgneal@17 289
bgneal@18 290 # Update root
bgneal@18 291 self.root = head.link[1]
bgneal@18 292 self.root.color = BLACK
bgneal@17 293
bgneal@19 294 return target, insert_flag
bgneal@19 295
bgneal@19 296 def remove(self, key):
bgneal@19 297 """Remove the given key from the tree.
bgneal@19 298
bgneal@19 299 Returns True if the key was found and removed and False if the key was
bgneal@19 300 not found in the tree.
bgneal@19 301
bgneal@19 302 """
bgneal@19 303 remove_flag = False # return value
bgneal@19 304
bgneal@19 305 if self.root is not None:
bgneal@19 306 # False/dummy tree root
bgneal@19 307 head = Node(key=None, value=None, color=BLACK)
bgneal@19 308 f = None # found item
bgneal@19 309 d = RIGHT # direction
bgneal@19 310
bgneal@19 311 # Set up helpers
bgneal@19 312 q = head
bgneal@19 313 g = p = None
bgneal@19 314 q.link[d] = self.root
bgneal@19 315
bgneal@19 316 # Search and push a red down to fix red violations as we go
bgneal@19 317 while q.link[d]:
bgneal@19 318 last = d
bgneal@19 319
bgneal@19 320 # Move the helpers down
bgneal@19 321 g = p
bgneal@19 322 p = q
bgneal@19 323 q = q.link[d]
bgneal@19 324 d = q.key < key
bgneal@19 325
bgneal@19 326 # Save the node with the matching data and keep going; we'll do
bgneal@19 327 # removal tasks at the end
bgneal@19 328 if q.key == key:
bgneal@19 329 f = q
bgneal@19 330
bgneal@19 331 # Push the red node down with rotations and color flips
bgneal@19 332 if not is_red(q) and not is_red(q.link[d]):
bgneal@19 333 if is_red(q.link[not d]):
bgneal@19 334 p.link[last] = self._single_rotate(q, d)
bgneal@19 335 p = p.link[last]
bgneal@19 336 elif not is_red(q.link[not d]):
bgneal@19 337 s = p.link[not last]
bgneal@19 338
bgneal@19 339 if s:
bgneal@19 340 if not is_red(s.link[not last]) and not is_red(s.link[last]):
bgneal@19 341 # Color flip
bgneal@19 342 p.color = BLACK
bgneal@19 343 s.color = RED
bgneal@19 344 q.color = RED
bgneal@19 345 else:
bgneal@19 346 d2 = g.link[1] is p
bgneal@19 347
bgneal@19 348 if is_red(s.link[last]):
bgneal@19 349 g.link[d2] = self._double_rotate(p, last)
bgneal@19 350 elif is_red(s.link[not last]):
bgneal@19 351 g.link[d2] = self._single_rotate(p, last)
bgneal@19 352
bgneal@19 353 # Ensure correct coloring
bgneal@19 354 q.color = g.link[d2].color = RED
bgneal@19 355 g.link[d2].link[0].color = BLACK
bgneal@19 356 g.link[d2].link[1].color = BLACK
bgneal@19 357
bgneal@19 358 # Replace and remove the saved node
bgneal@19 359 if f:
bgneal@19 360 f.key, f.value = q.key, q.value
bgneal@19 361 p.link[p.link[1] is q] = q.link[q.link[0] is None]
bgneal@19 362 q.free()
bgneal@19 363 self._size -= 1
bgneal@19 364 remove_flag = True
bgneal@19 365
bgneal@19 366 # Update root and make it black
bgneal@19 367 self.root = head.link[1]
bgneal@19 368 if self.root:
bgneal@19 369 self.root.color = BLACK
bgneal@19 370
bgneal@19 371 return remove_flag
bgneal@19 372
bgneal@20 373 def find(self, key):
bgneal@20 374 """Looks up the key in the tree and returns the corresponding value, or
bgneal@20 375 raises a KeyError if it does not exist in the tree.
bgneal@20 376
bgneal@20 377 """
bgneal@20 378 p = self.root
bgneal@20 379 while p:
bgneal@20 380 if p.key == key:
bgneal@20 381 return p.value
bgneal@20 382 else:
bgneal@20 383 d = p.key < key
bgneal@20 384 p = p.link[d]
bgneal@20 385
bgneal@20 386 raise KeyError
bgneal@20 387
bgneal@17 388
bgneal@17 389 if __name__ == '__main__':
bgneal@17 390 import random
bgneal@17 391
bgneal@18 392 for n in range(30):
bgneal@18 393 tree = Tree()
bgneal@18 394 tree.validate()
bgneal@17 395
bgneal@19 396 vals = []
bgneal@19 397 vals_set = set()
bgneal@18 398 for i in range(25):
bgneal@18 399 val = random.randint(0, 100)
bgneal@19 400 while val in vals_set:
bgneal@19 401 val = random.randint(0, 100)
bgneal@19 402 vals.append(val)
bgneal@19 403 vals_set.add(val)
bgneal@19 404
bgneal@18 405 tree.insert(val)
bgneal@19 406 tree.validate()
bgneal@17 407
bgneal@19 408 print 'Inserted in this order:', vals
bgneal@19 409 assert len(vals) == len(vals_set)
bgneal@19 410 keys = []
bgneal@18 411 for n in tree:
bgneal@18 412 print n.key,
bgneal@19 413 keys.append(n.key)
bgneal@19 414 print ' - len:', len(tree)
bgneal@18 415
bgneal@19 416 # delete in a random order
bgneal@19 417 random.shuffle(keys)
bgneal@19 418
bgneal@19 419 for k in keys:
bgneal@19 420 print 'Removing', k
bgneal@19 421 tree.remove(k)
bgneal@19 422 for n in tree:
bgneal@19 423 print n.key,
bgneal@19 424 print ' - len:', len(tree)
bgneal@19 425 tree.validate()