Package flumotion :: Package common :: Module avltree
[hide private]

Source Code for Module flumotion.common.avltree

  1  # -*- Mode: Python; test-case-name: flumotion.test.test_common_messages -*- 
  2  # vi:si:et:sw=4:sts=4:ts=4 
  3   
  4  # Flumotion - a streaming media server 
  5  # Copyright (C) 2004,2005,2006,2007,2008,2009 Fluendo, S.L. 
  6  # Copyright (C) 2010,2011 Flumotion Services, S.A. 
  7  # All rights reserved. 
  8  # 
  9  # This file may be distributed and/or modified under the terms of 
 10  # the GNU Lesser General Public License version 2.1 as published by 
 11  # the Free Software Foundation. 
 12  # This file is distributed without any warranty; without even the implied 
 13  # warranty of merchantability or fitness for a particular purpose. 
 14  # See "LICENSE.LGPL" in the source distribution for more information. 
 15  # 
 16  # Headers in this file shall remain intact. 
 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
34 -def height(tree):
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
47 -def debug(tree, level=0):
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
62 -def fromseq(seq):
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
70 -def _balance(hdiff, l, v, r, b):
71 """Internal method to rebalance an AVL tree, called as needed.""" 72 # if we have to rebalance, in the end the node has balance 0; 73 # for details see GNU libavl docs 74 if b < -1: 75 # rotate right 76 ll, lv, lr, lb = l 77 if lb == -1: 78 # easy case, lv is new root 79 new = node(ll, lv, node(lr, v, r, 0), 0) 80 if hdiff <= 0: 81 # deletion; maybe we decreased in height 82 old = node(l, v, r, b) 83 hdiff += height(new) - height(old) 84 else: 85 # we know that for insertion we don't increase in height 86 hdiff = 0 87 return hdiff, new 88 elif lb == 0: 89 # can only happen in deletion 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: # lb == +1 95 # lrv will be the new root 96 lrl, lrv, lrr, lrb = lr 97 if lrb == 0: # lr is the new node 98 newleftb = newrightb = 0 99 elif lrb == -1: 100 newleftb = 0 101 newrightb = +1 102 else: # lrb == +1 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 # deletion; maybe we decreased in height 109 old = node(l, v, r, b) 110 hdiff += height(new) - height(old) 111 else: 112 # we know that for insertion we don't increase in height 113 hdiff = 0 114 115 return hdiff, new 116 elif b > 1: 117 # rotate left 118 rl, rv, rr, rb = r 119 if rb == +1: 120 # easy case, rv is new root 121 new = node(node(l, v, rl, 0), rv, rr, 0) 122 if hdiff <= 0: 123 # deletion; maybe we decreased in height 124 old = node(l, v, r, b) 125 hdiff += height(new) - height(old) 126 else: 127 # we know that for insertion we don't increase in height 128 hdiff = 0 129 return hdiff, new 130 elif rb == 0: 131 # can only happen in deletion 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: # rb == -1 137 # rlv will be the new root 138 rll, rlv, rlr, rlb = rl 139 if rlb == 0: # rl is the new node 140 newleftb = newrightb = 0 141 elif rlb == +1: 142 newleftb = -1 143 newrightb = 0 144 else: # rlb == -1 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 # deletion; maybe we decreased in height 151 old = node(l, v, r, b) 152 hdiff += height(new) - height(old) 153 else: 154 # we know that for insertion we don't increase in height 155 hdiff = 0 156 return hdiff, new 157 else: 158 return hdiff, node(l, v, r, b)
159 160
161 -def insert(tree, value):
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
186 -def delete(tree, value):
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 # overall height only changes if left was taller before 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 # overall height only changes if left was 213 # taller before 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 # overall height only changes if right was 222 # taller before 223 hdiff = 0 224 b -= 1 225 return _balance(hdiff, l, v, newr, b) 226 else: 227 # we have found the node! 228 if r is None: 229 # no right link, just replace with left 230 return -1, l 231 else: 232 newv, hdiff, newr = popmin(r) 233 if hdiff != 0: 234 if b <= 0: 235 # overall height only changes if right was 236 # taller before 237 hdiff = 0 238 b -= 1 239 return _balance(hdiff, l, newv, newr, b) 240 241
242 -def lookup(tree, value):
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
257 -def iterate(tree):
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
269 -def iteratereversed(tree):
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
281 -class AVLTree(object):
282
283 - def __init__(self, seq=()):
284 self._len = len(seq) 285 self.tree = fromseq(seq)
286
287 - def insert(self, value):
288 _, self.tree = insert(self.tree, value) 289 self._len += 1
290
291 - def delete(self, value):
292 _, self.tree = delete(self.tree, value) 293 self._len -= 1
294
295 - def __contains__(self, value):
296 return bool(lookup(self.tree, value))
297
298 - def __len__(self):
299 return self._len
300
301 - def __iter__(self):
302 return iterate(self.tree)
303
304 - def iterreversed(self):
305 return iteratereversed(self.tree)
306