Mercurial > public > think_complexity

annotate redblacktree.py @ 24:5c2c4ce095ef

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
A stab at the L(p)/L(0) plot. I still don't quite get how the graphs in the Watts and Strogatz paper were generated. My results have basically the same shape, but don't converge to 0. I'm not sure how this is possible if the rewire function does not remove edges.
author Brian Neal <bgneal@gmail.com>
date Thu, 03 Jan 2013 18:41:13 -0600
parents 0326803882ad
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()