Mercurial > public > think_complexity

comparison 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
comparison
equal deleted inserted replaced
16:a00e97bcdb4a 17:977628018b4b
1 """A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think 1 """A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think
2 Complexity_ book. 2 Complexity_ book.
3 3
4 http://greenteapress.com/complexity 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
5 11
6 Copyright (C) 2012 Brian G. Neal. 12 Copyright (C) 2012 Brian G. Neal.
7 13
8 Permission is hereby granted, free of charge, to any person obtaining a copy of 14 Permission is hereby granted, free of charge, to any person obtaining a copy of
9 this software and associated documentation files (the "Software"), to deal in 15 this software and associated documentation files (the "Software"), to deal in
62 """ 68 """
63 g = self.grandparent 69 g = self.grandparent
64 if g: 70 if g:
65 return g.left if g.right is self.parent else g.right 71 return g.left if g.right is self.parent else g.right
66 return None 72 return None
73
74 def __str__(self):
75 c = 'B' if self.color == BLACK else 'R'
76 if self.value:
77 return '({}: {} => {})'.format(c, self.key, self.value)
78 else:
79 return '({}: {})'.format(c, self.key)
80
81
82 class Tree(object):
83 """A red-black Tree class.
84
85 A red-black tree is a binary search tree with the following properties:
86
87 1. A node is either red or black.
88 2. The root is black.
89 3. All leaves are black.
90 4. Both children of every red node are black.
91 5. Every simple path from a given node to any descendant leaf contains
92 the same number of black nodes.
93
94 These rules ensure that the path from the root to the furthest leaf is no
95 more than twice as long as the path from the root to the nearest leaf. Thus
96 the tree is roughly height-balanced.
97
98 """
99
100 def __init__(self):
101 self.root = None
102
103 def __iter__(self):
104 return self._inorder(self.root)
105
106 def _inorder(self, node):
107 """A generator to perform an inorder traversal of the nodes in the
108 tree starting at the given node.
109
110 """
111 if node.left:
112 for n in self._inorder(node.left):
113 yield n
114
115 yield node
116
117 if node.right:
118 for n in self._inorder(node.right):
119 yield n
120
121 def insert(self, key, value=None):
122 node = Node(key=key, value=value, color=RED)
123
124 # trivial case of inserting into an empty tree:
125 if self.root is None:
126 self.root = node
127 self.root.color = BLACK
128 return
129
130 # Find a spot to insert the new red node:
131
132 x = self.root
133 while x is not None:
134 p = x
135 if key < x.key:
136 x = x.left
137 else:
138 x = x.right
139
140 node.parent = p # p is the new node's parent
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:
152 p = node.parent
153
154 # Case 1: root node
155 if p is None:
156 node.color = BLACK
157 break
158
159 # Case 2: parent is black
160 if p.color == BLACK:
161 break
162
163 # Case 3: parent and uncle are red
164 u = node.uncle
165 if u and u.color == RED:
166 # repaint parent & uncle black; grandparent becomes red:
167 p.color = BLACK
168 u.color = BLACK
169 gp = node.grandparent
170 gp.color = RED
171
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
268
269 if __name__ == '__main__':
270 import random
271
272 tree = Tree()
273
274 for i in range(20):
275 val = random.randint(0, 100)
276 tree.insert(val)
277
278 for n in tree:
279 print n.key,