Mercurial > public > think_complexity

annotate redblacktree.py @ 17:977628018b4b

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
The insert operation on the red-black tree seems to be working.
author Brian Neal <bgneal@gmail.com>
date Wed, 19 Dec 2012 22:29:09 -0600
parents a00e97bcdb4a
children 92e2879e2e33
rev   line source
bgneal@15 1 """A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think
bgneal@15 2 Complexity_ book.
bgneal@15 3
bgneal@15 4 http://greenteapress.com/complexity
bgneal@15 5
bgneal@17 6 This code is based on the description of the red-black tree at
bgneal@17 7 http://en.wikipedia.org/wiki/Red-black_tree.
bgneal@17 8
bgneal@17 9 Some code and ideas were taken from code by Darren Hart at
bgneal@17 10 http://dvhart.com/darren/files/rbtree.py
bgneal@17 11
bgneal@15 12 Copyright (C) 2012 Brian G. Neal.
bgneal@15 13
bgneal@15 14 Permission is hereby granted, free of charge, to any person obtaining a copy of
bgneal@15 15 this software and associated documentation files (the "Software"), to deal in
bgneal@15 16 the Software without restriction, including without limitation the rights to
bgneal@15 17 use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
bgneal@15 18 the Software, and to permit persons to whom the Software is furnished to do so,
bgneal@15 19 subject to the following conditions:
bgneal@15 20
bgneal@15 21 The above copyright notice and this permission notice shall be included in all
bgneal@15 22 copies or substantial portions of the Software.
bgneal@15 23
bgneal@15 24 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
bgneal@15 25 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
bgneal@15 26 FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
bgneal@15 27 COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
bgneal@15 28 IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
bgneal@15 29 CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
bgneal@15 30
bgneal@15 31 """
bgneal@15 32
bgneal@16 33 BLACK, RED = range(2)
bgneal@15 34
bgneal@15 35 class Node(object):
bgneal@15 36 """A node class for red-black trees.
bgneal@15 37
bgneal@15 38 A node has an optional parent, and optional left and right children. Each
bgneal@15 39 node also has a color, either red or black. A node has a key and an optional
bgneal@15 40 value. The key is used to order the red black tree by calling the "<"
bgneal@16 41 operator when comparing keys. The optional value is useful for using the
bgneal@16 42 red-black tree to implement a map datastructure.
bgneal@15 43
bgneal@15 44 In a red-black tree, nil children are always considered black.
bgneal@15 45
bgneal@15 46 """
bgneal@15 47
bgneal@15 48 def __init__(self, key, color=RED, value=None):
bgneal@15 49 self.key = key
bgneal@15 50 self.value = value
bgneal@15 51 self.color = color
bgneal@15 52
bgneal@15 53 self.parent = None
bgneal@15 54 self.left = None
bgneal@15 55 self.right = None
bgneal@15 56
bgneal@15 57 @property
bgneal@15 58 def grandparent(self):
bgneal@15 59 """Return the grandparent of this node, or None if it does not exist."""
bgneal@15 60 return self.parent.parent if self.parent else None
bgneal@15 61
bgneal@15 62 @property
bgneal@15 63 def uncle(self):
bgneal@15 64 """Return this node's uncle if it exists, or None if not.
bgneal@15 65
bgneal@15 66 An uncle is a parent's sibling.
bgneal@15 67
bgneal@15 68 """
bgneal@15 69 g = self.grandparent
bgneal@15 70 if g:
bgneal@15 71 return g.left if g.right is self.parent else g.right
bgneal@15 72 return 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@17 82 class Tree(object):
bgneal@17 83 """A red-black Tree class.
bgneal@17 84
bgneal@17 85 A red-black tree is a binary search tree with the following properties:
bgneal@17 86
bgneal@17 87 1. A node is either red or black.
bgneal@17 88 2. The root is black.
bgneal@17 89 3. All leaves are black.
bgneal@17 90 4. Both children of every red node are black.
bgneal@17 91 5. Every simple path from a given node to any descendant leaf contains
bgneal@17 92 the same number of black nodes.
bgneal@17 93
bgneal@17 94 These rules ensure that the path from the root to the furthest leaf is no
bgneal@17 95 more than twice as long as the path from the root to the nearest leaf. Thus
bgneal@17 96 the tree is roughly height-balanced.
bgneal@17 97
bgneal@17 98 """
bgneal@17 99
bgneal@17 100 def __init__(self):
bgneal@17 101 self.root = None
bgneal@17 102
bgneal@17 103 def __iter__(self):
bgneal@17 104 return self._inorder(self.root)
bgneal@17 105
bgneal@17 106 def _inorder(self, node):
bgneal@17 107 """A generator to perform an inorder traversal of the nodes in the
bgneal@17 108 tree starting at the given node.
bgneal@17 109
bgneal@17 110 """
bgneal@17 111 if node.left:
bgneal@17 112 for n in self._inorder(node.left):
bgneal@17 113 yield n
bgneal@17 114
bgneal@17 115 yield node
bgneal@17 116
bgneal@17 117 if node.right:
bgneal@17 118 for n in self._inorder(node.right):
bgneal@17 119 yield n
bgneal@17 120
bgneal@17 121 def insert(self, key, value=None):
bgneal@17 122 node = Node(key=key, value=value, color=RED)
bgneal@17 123
bgneal@17 124 # trivial case of inserting into an empty tree:
bgneal@17 125 if self.root is None:
bgneal@17 126 self.root = node
bgneal@17 127 self.root.color = BLACK
bgneal@17 128 return
bgneal@17 129
bgneal@17 130 # Find a spot to insert the new red node:
bgneal@17 131
bgneal@17 132 x = self.root
bgneal@17 133 while x is not None:
bgneal@17 134 p = x
bgneal@17 135 if key < x.key:
bgneal@17 136 x = x.left
bgneal@17 137 else:
bgneal@17 138 x = x.right
bgneal@17 139
bgneal@17 140 node.parent = p # p is the new node's parent
bgneal@17 141
bgneal@17 142 # p now has 1 child; decide if the new node goes on the left or right
bgneal@17 143
bgneal@17 144 if key < p.key:
bgneal@17 145 p.left = node
bgneal@17 146 else:
bgneal@17 147 p.right = node
bgneal@17 148
bgneal@17 149 # ensure the new tree follows the red-black rules:
bgneal@17 150
bgneal@17 151 while True:
bgneal@17 152 p = node.parent
bgneal@17 153
bgneal@17 154 # Case 1: root node
bgneal@17 155 if p is None:
bgneal@17 156 node.color = BLACK
bgneal@17 157 break
bgneal@17 158
bgneal@17 159 # Case 2: parent is black
bgneal@17 160 if p.color == BLACK:
bgneal@17 161 break
bgneal@17 162
bgneal@17 163 # Case 3: parent and uncle are red
bgneal@17 164 u = node.uncle
bgneal@17 165 if u and u.color == RED:
bgneal@17 166 # repaint parent & uncle black; grandparent becomes red:
bgneal@17 167 p.color = BLACK
bgneal@17 168 u.color = BLACK
bgneal@17 169 gp = node.grandparent
bgneal@17 170 gp.color = RED
bgneal@17 171
bgneal@17 172 # enforce the rules at the grandparent level
bgneal@17 173 node = gp
bgneal@17 174 continue
bgneal@17 175
bgneal@17 176 # gp is the grandparent; it can't be None because we know the parent
bgneal@17 177 # is red at this point
bgneal@17 178 gp = p.parent
bgneal@17 179
bgneal@17 180 # Case 4: At this point we know the parent is red and the uncle is
bgneal@17 181 # black.
bgneal@17 182 # If the new node is a right child of the parent, and the
bgneal@17 183 # parent is on the left of the grandparent => left rotation on the
bgneal@17 184 # parent.
bgneal@17 185 # If the new node is a left child of the parent, and the parent is
bgneal@17 186 # on the right of the grandparent => right rotation on the parent.
bgneal@17 187
bgneal@17 188 if node is p.right and p is gp.left:
bgneal@17 189 self._rotate_left(p)
bgneal@17 190 node = node.left
bgneal@17 191 elif node is p.left and p is gp.right:
bgneal@17 192 self._rotate_right(p)
bgneal@17 193 node = node.right
bgneal@17 194
bgneal@17 195 # Note that case 4 must fall through to case 5 to fix up the former parent
bgneal@17 196 # node, which is now the child of the new red node.
bgneal@17 197
bgneal@17 198 # Case 5: The parent is red and the uncle is black.
bgneal@17 199 # If the new node is the left child of the parent and the parent is
bgneal@17 200 # the left child of the grandparent => rotate right on the
bgneal@17 201 # grandparent.
bgneal@17 202 # If the new node is the right child of the parent and the parent
bgneal@17 203 # is the right child of the grandparent => rotate left on the
bgneal@17 204 # grandparent.
bgneal@17 205
bgneal@17 206 p = node.parent
bgneal@17 207 gp = p.parent
bgneal@17 208
bgneal@17 209 p.color = BLACK
bgneal@17 210 gp.color = RED
bgneal@17 211
bgneal@17 212 if node is p.left and p is gp.left:
bgneal@17 213 self._rotate_right(gp)
bgneal@17 214 elif node is p.right and p is gp.right: #TODO: can this be an else?
bgneal@17 215 self._rotate_left(gp)
bgneal@17 216 break
bgneal@17 217
bgneal@17 218 def _rotate_right(self, node):
bgneal@17 219 """Rotate the tree right on the given node."""
bgneal@17 220
bgneal@17 221 p = node.parent
bgneal@17 222 left = node.left
bgneal@17 223
bgneal@17 224 if p:
bgneal@17 225 if node is p.right:
bgneal@17 226 p.right = left
bgneal@17 227 else:
bgneal@17 228 p.left = left
bgneal@17 229
bgneal@17 230 left.parent = p
bgneal@17 231 node.left = left.right
bgneal@17 232
bgneal@17 233 if node.left:
bgneal@17 234 node.left.parent = node
bgneal@17 235
bgneal@17 236 left.right = node
bgneal@17 237 node.parent = left
bgneal@17 238
bgneal@17 239 # fix up the root if necessary
bgneal@17 240 if self.root is node:
bgneal@17 241 self.root = left
bgneal@17 242
bgneal@17 243 def _rotate_left(self, node):
bgneal@17 244 """Rotate the tree left on the given node."""
bgneal@17 245
bgneal@17 246 p = node.parent
bgneal@17 247 right = node.right
bgneal@17 248
bgneal@17 249 if p:
bgneal@17 250 if node is p.right:
bgneal@17 251 p.right = right
bgneal@17 252 else:
bgneal@17 253 p.left = right
bgneal@17 254
bgneal@17 255 right.parent = p
bgneal@17 256 node.right = right.left
bgneal@17 257
bgneal@17 258 if node.right:
bgneal@17 259 node.right.parent = node
bgneal@17 260
bgneal@17 261 right.left = node
bgneal@17 262 node.parent = right
bgneal@17 263
bgneal@17 264 # fix up the root if necessary
bgneal@17 265 if self.root is node:
bgneal@17 266 self.root = right
bgneal@17 267
bgneal@17 268
bgneal@17 269 if __name__ == '__main__':
bgneal@17 270 import random
bgneal@17 271
bgneal@17 272 tree = Tree()
bgneal@17 273
bgneal@17 274 for i in range(20):
bgneal@17 275 val = random.randint(0, 100)
bgneal@17 276 tree.insert(val)
bgneal@17 277
bgneal@17 278 for n in tree:
bgneal@17 279 print n.key,