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