Mercurial > public > think_complexity

comparison 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
comparison
equal deleted inserted replaced
17:977628018b4b 18:92e2879e2e33
1 """A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think 1 """
2 Complexity_ book.
3
4 http://greenteapress.com/complexity
5
6 This code is based on the description of the red-black tree at
7 http://en.wikipedia.org/wiki/Red-black_tree.
8
9 Some code and ideas were taken from code by Darren Hart at
10 http://dvhart.com/darren/files/rbtree.py
11
12 Copyright (C) 2012 Brian G. Neal. 2 Copyright (C) 2012 Brian G. Neal.
13 3
14 Permission is hereby granted, free of charge, to any person obtaining a copy of 4 Permission is hereby granted, free of charge, to any person obtaining a copy of
15 this software and associated documentation files (the "Software"), to deal in 5 this software and associated documentation files (the "Software"), to deal in
16 the Software without restriction, including without limitation the rights to 6 the Software without restriction, including without limitation the rights to
26 FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 16 FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
27 COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 17 COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
28 IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 18 IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
29 CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 19 CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
30 20
21 ----
22
23 A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think
24 Complexity_ book.
25
26 http://greenteapress.com/complexity
27
28 Red black trees are described on Wikipedia:
29 http://en.wikipedia.org/wiki/Red-black_tree.
30
31 This is basically a Python implementation of Julienne Walker's Red Black Trees
32 tutorial found at:
33 http://www.eternallyconfuzzled.com/tuts/datastructures/jsw_tut_rbtree.aspx
34 We implement Julienne's top-down insertion and deletion algorithms here.
35
36 Some ideas were also borrowed from code by Darren Hart at
37 http://dvhart.com/darren/files/rbtree.py
38
31 """ 39 """
32 40
33 BLACK, RED = range(2) 41 BLACK, RED = range(2)
42 LEFT, RIGHT = range(2)
43
44
45 class TreeError(Exception):
46 """Base exception class for red-black tree errors."""
47 pass
48
34 49
35 class Node(object): 50 class Node(object):
36 """A node class for red-black trees. 51 """A node class for red-black trees.
37 52
38 A node has an optional parent, and optional left and right children. Each 53 * A node has a color, either RED or BLACK.
39 node also has a color, either red or black. A node has a key and an optional 54 * A node has a key and an optional value.
40 value. The key is used to order the red black tree by calling the "<" 55
41 operator when comparing keys. The optional value is useful for using the 56 * The key is used to order the red-black tree by calling the "<"
42 red-black tree to implement a map datastructure. 57 operator when comparing two nodes.
43 58 * The value is useful for using the red-black tree to implement a map
44 In a red-black tree, nil children are always considered black. 59 datastructure.
60
61 * Nodes have exactly 2 link fields which we represent as a list of
62 2 elements to represent the left and right children of the node. This list
63 representation was borrowed from Julienne Walker as it simplifies some of
64 the code. Element 0 is the LEFT child and element 1 is the RIGHT child.
45 65
46 """ 66 """
47 67
48 def __init__(self, key, color=RED, value=None): 68 def __init__(self, key, value=None, color=RED):
49 self.key = key 69 self.key = key
50 self.value = value 70 self.value = value
51 self.color = color 71 self.color = color
52 72 self.link = [None, None]
53 self.parent = None
54 self.left = None
55 self.right = None
56
57 @property
58 def grandparent(self):
59 """Return the grandparent of this node, or None if it does not exist."""
60 return self.parent.parent if self.parent else None
61
62 @property
63 def uncle(self):
64 """Return this node's uncle if it exists, or None if not.
65
66 An uncle is a parent's sibling.
67
68 """
69 g = self.grandparent
70 if g:
71 return g.left if g.right is self.parent else g.right
72 return None
73 73
74 def __str__(self): 74 def __str__(self):
75 c = 'B' if self.color == BLACK else 'R' 75 c = 'B' if self.color == BLACK else 'R'
76 if self.value: 76 if self.value:
77 return '({}: {} => {})'.format(c, self.key, self.value) 77 return '({}: {} => {})'.format(c, self.key, self.value)
78 else: 78 else:
79 return '({}: {})'.format(c, self.key) 79 return '({}: {})'.format(c, self.key)
80 80
81 81
82 def is_red(n):
83 """Return True if the given Node n is RED and False otherwise.
84
85 If the node is None, then it is considered to be BLACK, and False is
86 returned.
87
88 """
89 return n is not None and n.color == RED
90
91
82 class Tree(object): 92 class Tree(object):
83 """A red-black Tree class. 93 """A red-black Tree class.
84 94
85 A red-black tree is a binary search tree with the following properties: 95 A red-black tree is a binary search tree with the following properties:
86 96
106 def _inorder(self, node): 116 def _inorder(self, node):
107 """A generator to perform an inorder traversal of the nodes in the 117 """A generator to perform an inorder traversal of the nodes in the
108 tree starting at the given node. 118 tree starting at the given node.
109 119
110 """ 120 """
111 if node.left: 121 if node.link[0]:
112 for n in self._inorder(node.left): 122 for n in self._inorder(node.link[0]):
113 yield n 123 yield n
114 124
115 yield node 125 yield node
116 126
117 if node.right: 127 if node.link[1]:
118 for n in self._inorder(node.right): 128 for n in self._inorder(node.link[1]):
119 yield n 129 yield n
120 130
131 def _single_rotate(self, root, d):
132 """Perform a single rotation about the node 'root' in the given
133 direction 'd' (LEFT or RIGHT).
134
135 The old root is set to RED and the new root is set to BLACK.
136
137 Returns the new root.
138
139 """
140 nd = not d
141 save = root.link[nd]
142
143 root.link[nd] = save.link[d]
144 save.link[d] = root
145
146 root.color = RED
147 save.color = BLACK
148
149 return save
150
151 def _double_rotate(self, root, d):
152 """Perform two single rotations about the node root in the direction d.
153
154 The new root is returned.
155
156 """
157 nd = not d
158 root.link[nd] = self._single_rotate(root.link[nd], nd)
159 return self._single_rotate(root, d)
160
161 def validate(self, root=None):
162 """Checks to see if the red-black tree validates at the given node, i.e.
163 all red-black tree rules are valid. If root is None, the root of the
164 tree is used as the starting point.
165
166 If any rules are violated, a TreeError is raised.
167
168 Returns the black height of the tree rooted at root.
169
170 """
171 if root is None:
172 root = self.root
173
174 return self._validate(root)
175
176 def _validate(self, root):
177 """Internal implementation of the validate() method."""
178
179 if root is None:
180 return 1
181
182 ln = root.link[0]
183 rn = root.link[1]
184
185 # Check for consecutive red links
186
187 if is_red(root) and (is_red(ln) or is_red(rn)):
188 raise TreeError('red violation')
189
190 lh = self._validate(ln)
191 rh = self._validate(rn)
192
193 # Check for invalid binary search tree
194
195 if (ln and ln.key >= root.key) or (rn and rn.key <= root.key):
196 raise TreeError('binary tree violation')
197
198 # Check for black height mismatch
199 if lh != rh:
200 raise TreeError('black violation')
201
202 # Only count black links
203
204 return lh if is_red(root) else lh + 1
205
121 def insert(self, key, value=None): 206 def insert(self, key, value=None):
122 node = Node(key=key, value=value, color=RED) 207 """Insert the (key, value) pair into the tree."""
123 208
124 # trivial case of inserting into an empty tree: 209 # Check for the empty tree case:
125 if self.root is None: 210 if self.root is None:
126 self.root = node 211 self.root = Node(key=key, value=value, color=BLACK)
127 self.root.color = BLACK
128 return 212 return
129 213
130 # Find a spot to insert the new red node: 214 # False/dummy tree root
131 215 head = Node(key=None, value=None, color=BLACK)
132 x = self.root 216 d = LEFT
133 while x is not None: 217 last = LEFT
134 p = x 218
135 if key < x.key: 219 t = head
136 x = x.left 220 g = p = None
137 else: 221 t.link[1] = self.root
138 x = x.right 222 q = self.root
139 223
140 node.parent = p # p is the new node's parent 224 # Search down the tree
141
142 # p now has 1 child; decide if the new node goes on the left or right
143
144 if key < p.key:
145 p.left = node
146 else:
147 p.right = node
148
149 # ensure the new tree follows the red-black rules:
150
151 while True: 225 while True:
152 p = node.parent 226 if q is None:
153 227 # Insert new node at the bottom
154 # Case 1: root node 228 p.link[d] = q = Node(key=key, value=value, color=RED)
155 if p is None: 229 elif is_red(q.link[0]) and is_red(q.link[1]):
156 node.color = BLACK 230 # Color flip
231 q.color = RED
232 q.link[0].color = BLACK
233 q.link[1].color = BLACK
234
235 # Fix red violation
236 if is_red(q) and is_red(p):
237 d2 = t.link[1] is g
238
239 if q is p.link[last]:
240 t.link[d2] = self._single_rotate(g, not last)
241 else:
242 t.link[d2] = self._double_rotate(g, not last)
243
244 # Stop if found
245 if q.key == key:
157 break 246 break
158 247
159 # Case 2: parent is black 248 last = d
160 if p.color == BLACK: 249 d = q.key < key
161 break 250
162 251 # Update helpers
163 # Case 3: parent and uncle are red 252 if g is not None:
164 u = node.uncle 253 t = g;
165 if u and u.color == RED: 254 g = p
166 # repaint parent & uncle black; grandparent becomes red: 255 p = q
167 p.color = BLACK 256 q = q.link[d]
168 u.color = BLACK 257
169 gp = node.grandparent 258 # Update root
170 gp.color = RED 259 self.root = head.link[1]
171 260 self.root.color = BLACK
172 # enforce the rules at the grandparent level
173 node = gp
174 continue
175
176 # gp is the grandparent; it can't be None because we know the parent
177 # is red at this point
178 gp = p.parent
179
180 # Case 4: At this point we know the parent is red and the uncle is
181 # black.
182 # If the new node is a right child of the parent, and the
183 # parent is on the left of the grandparent => left rotation on the
184 # parent.
185 # If the new node is a left child of the parent, and the parent is
186 # on the right of the grandparent => right rotation on the parent.
187
188 if node is p.right and p is gp.left:
189 self._rotate_left(p)
190 node = node.left
191 elif node is p.left and p is gp.right:
192 self._rotate_right(p)
193 node = node.right
194
195 # Note that case 4 must fall through to case 5 to fix up the former parent
196 # node, which is now the child of the new red node.
197
198 # Case 5: The parent is red and the uncle is black.
199 # If the new node is the left child of the parent and the parent is
200 # the left child of the grandparent => rotate right on the
201 # grandparent.
202 # If the new node is the right child of the parent and the parent
203 # is the right child of the grandparent => rotate left on the
204 # grandparent.
205
206 p = node.parent
207 gp = p.parent
208
209 p.color = BLACK
210 gp.color = RED
211
212 if node is p.left and p is gp.left:
213 self._rotate_right(gp)
214 elif node is p.right and p is gp.right: #TODO: can this be an else?
215 self._rotate_left(gp)
216 break
217
218 def _rotate_right(self, node):
219 """Rotate the tree right on the given node."""
220
221 p = node.parent
222 left = node.left
223
224 if p:
225 if node is p.right:
226 p.right = left
227 else:
228 p.left = left
229
230 left.parent = p
231 node.left = left.right
232
233 if node.left:
234 node.left.parent = node
235
236 left.right = node
237 node.parent = left
238
239 # fix up the root if necessary
240 if self.root is node:
241 self.root = left
242
243 def _rotate_left(self, node):
244 """Rotate the tree left on the given node."""
245
246 p = node.parent
247 right = node.right
248
249 if p:
250 if node is p.right:
251 p.right = right
252 else:
253 p.left = right
254
255 right.parent = p
256 node.right = right.left
257
258 if node.right:
259 node.right.parent = node
260
261 right.left = node
262 node.parent = right
263
264 # fix up the root if necessary
265 if self.root is node:
266 self.root = right
267 261
268 262
269 if __name__ == '__main__': 263 if __name__ == '__main__':
270 import random 264 import random
271 265
272 tree = Tree() 266 for n in range(30):
273 267 tree = Tree()
274 for i in range(20): 268 tree.validate()
275 val = random.randint(0, 100) 269
276 tree.insert(val) 270 for i in range(25):
277 271 val = random.randint(0, 100)
278 for n in tree: 272 tree.insert(val)
279 print n.key, 273
274 for n in tree:
275 print n.key,
276 print
277
278 tree.validate()