1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22 """self-balancing binary search tree.
23 A pure python functional-style self-balancing binary search tree
24 implementation, with an object-oriented wrapper. Useful for maintaining
25 sorted sets, traversing sets in order, and closest-match lookups.
26 """
27
28 __version__ = "$Rev$"
29
30
31 -def node(l, v, r, b):
32 """Make an AVL tree node, consisting of a left tree, a value, a
33 right tree, and the "balance factor": the difference in lengths
34 between the right and left sides, respectively."""
35 return (l, v, r, b)
36
37
39 """Return the height of an AVL tree. Relies on the balance factors
40 being consistent."""
41 if tree is None:
42 return 0
43 else:
44 l, v, r, b = tree
45 if b <= 0:
46 return height(l) + 1
47 else:
48 return height(r) + 1
49
50
52 """Print out a debugging representation of an AVL tree."""
53 if tree is None:
54 return
55 l, v, r, b = tree
56 debug(l, level+1)
57 bchr = {-2: '--',
58 -1: '-',
59 0: '0',
60 1: '+',
61 2: '++'}.get(b, '?')
62 print '%s%s: %r' % (' '*level, bchr, v)
63 debug(r, level+1)
64
65
67 """Populate and return an AVL tree from an iterable sequence."""
68 t = None
69 for x in seq:
70 _, t = insert(t, x)
71 return t
72
73
75 """Internal method to rebalance an AVL tree, called as needed."""
76
77
78 if b < -1:
79
80 ll, lv, lr, lb = l
81 if lb == -1:
82
83 new = node(ll, lv, node(lr, v, r, 0), 0)
84 if hdiff <= 0:
85
86 old = node(l, v, r, b)
87 hdiff += height(new) - height(old)
88 else:
89
90 hdiff = 0
91 return hdiff, new
92 elif lb == 0:
93
94 new = node(ll, lv, node(lr, v, r, -1), +1)
95 old = node(l, v, r, b)
96 hdiff += height(new) - height(old)
97 return hdiff, new
98 else:
99
100 lrl, lrv, lrr, lrb = lr
101 if lrb == 0:
102 newleftb = newrightb = 0
103 elif lrb == -1:
104 newleftb = 0
105 newrightb = +1
106 else:
107 newleftb = -1
108 newrightb = 0
109 new = node(node(ll, lv, lrl, newleftb), lrv,
110 node(lrr, v, r, newrightb), 0)
111 if hdiff <= 0:
112
113 old = node(l, v, r, b)
114 hdiff += height(new) - height(old)
115 else:
116
117 hdiff = 0
118
119 return hdiff, new
120 elif b > 1:
121
122 rl, rv, rr, rb = r
123 if rb == +1:
124
125 new = node(node(l, v, rl, 0), rv, rr, 0)
126 if hdiff <= 0:
127
128 old = node(l, v, r, b)
129 hdiff += height(new) - height(old)
130 else:
131
132 hdiff = 0
133 return hdiff, new
134 elif rb == 0:
135
136 new = node(node(l, v, rl, +1), rv, rr, -1)
137 old = node(l, v, r, b)
138 hdiff += height(new) - height(old)
139 return hdiff, new
140 else:
141
142 rll, rlv, rlr, rlb = rl
143 if rlb == 0:
144 newleftb = newrightb = 0
145 elif rlb == +1:
146 newleftb = -1
147 newrightb = 0
148 else:
149 newleftb = 0
150 newrightb = +1
151 new = node(node(l, v, rll, newleftb), rlv,
152 node(rlr, rv, rr, newrightb), 0)
153 if hdiff <= 0:
154
155 old = node(l, v, r, b)
156 hdiff += height(new) - height(old)
157 else:
158
159 hdiff = 0
160 return hdiff, new
161 else:
162 return hdiff, node(l, v, r, b)
163
164
166 """Insert a value into an AVL tree. Returns a tuple of
167 (heightdifference, tree). The original tree is unmodified."""
168 if tree is None:
169 return 1, (None, value, None, 0)
170 else:
171 l, v, r, b = tree
172 if value < v:
173 hdiff, newl = insert(l, value)
174 if hdiff > 0:
175 if b > 0:
176 hdiff = 0
177 b -= 1
178 return _balance(hdiff, newl, v, r, b)
179 elif value > v:
180 hdiff, newr = insert(r, value)
181 if hdiff > 0:
182 if b < 0:
183 hdiff = 0
184 b += 1
185 return _balance(hdiff, l, v, newr, b)
186 else:
187 raise ValueError('tree already has value %r' % (value, ))
188
189
191 """Delete a value from an AVL tree. Like L{insert}, returns a tuple
192 of (heightdifference, tree). The original tree is unmodified."""
193
194 def popmin((l, v, r, b)):
195 if l is None:
196 minv = v
197 return minv, -1, r
198 else:
199 minv, hdiff, newl = popmin(l)
200 if hdiff != 0:
201 if b >= 0:
202
203 hdiff = 0
204 b += 1
205
206 return (minv, ) + _balance(hdiff, newl, v, r, b)
207
208 if tree is None:
209 raise ValueError('tree has no value %r' % (value, ))
210 else:
211 l, v, r, b = tree
212 if value < v:
213 hdiff, newl = delete(l, value)
214 if hdiff != 0:
215 if b >= 0:
216
217
218 hdiff = 0
219 b += 1
220 return _balance(hdiff, newl, v, r, b)
221 elif value > v:
222 hdiff, newr = delete(r, value)
223 if hdiff != 0:
224 if b <= 0:
225
226
227 hdiff = 0
228 b -= 1
229 return _balance(hdiff, l, v, newr, b)
230 else:
231
232 if r is None:
233
234 return -1, l
235 else:
236 newv, hdiff, newr = popmin(r)
237 if hdiff != 0:
238 if b <= 0:
239
240
241 hdiff = 0
242 b -= 1
243 return _balance(hdiff, l, newv, newr, b)
244
245
247 """Look up a node in an AVL tree. Returns a node tuple or False if
248 the value was not found."""
249 if tree is None:
250 return False
251 else:
252 l, v, r, b = tree
253 if value < v:
254 return lookup(l, v)
255 elif value > v:
256 return lookup(r, v)
257 else:
258 return tree
259
260
262 """Iterate over an AVL tree, starting with the lowest-ordered
263 value."""
264 if tree is not None:
265 l, v, r, b = tree
266 for x in iterate(l):
267 yield x
268 yield v
269 for x in iterate(r):
270 yield x
271
272
274 """Iterate over an AVL tree, starting with the highest-ordered
275 value."""
276 if tree is not None:
277 l, v, r, b = tree
278 for x in iteratereversed(r):
279 yield x
280 yield v
281 for x in iteratereversed(l):
282 yield x
283
284
286
288 self._len = len(seq)
289 self.tree = fromseq(seq)
290
292 _, self.tree = insert(self.tree, value)
293 self._len += 1
294
296 _, self.tree = delete(self.tree, value)
297 self._len -= 1
298
300 return bool(lookup(self.tree, value))
301
304
307
310