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