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) 2007 Fluendo, S.L. (www.fluendo.com). 
  6  # All rights reserved. 
  7   
  8  # This file may be distributed and/or modified under the terms of 
  9  # the GNU General Public License version 2 as published by 
 10  # the Free Software Foundation. 
 11  # This file is distributed without any warranty; without even the implied 
 12  # warranty of merchantability or fitness for a particular purpose. 
 13  # See "LICENSE.GPL" in the source distribution for more information. 
 14   
 15  # Licensees having purchased or holding a valid Flumotion Advanced 
 16  # Streaming Server license may use this file in accordance with the 
 17  # Flumotion Advanced Streaming Server Commercial License Agreement. 
 18  # See "LICENSE.Flumotion" in the source distribution for more information. 
 19   
 20  # Headers in this file shall remain intact. 
 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: 7138 $" 
 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
38 -def height(tree):
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
51 -def debug(tree, level=0):
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
66 -def fromseq(seq):
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
74 -def _balance(hdiff, l, v, r, b):
75 """Internal method to rebalance an AVL tree, called as needed.""" 76 # if we have to rebalance, in the end the node has balance 0; 77 # for details see GNU libavl docs 78 if b < -1: 79 # rotate right 80 ll, lv, lr, lb = l 81 if lb == -1: 82 # easy case, lv is new root 83 new = node(ll, lv, node(lr, v, r, 0), 0) 84 if hdiff <= 0: 85 # deletion; maybe we decreased in height 86 old = node(l, v, r, b) 87 hdiff += height(new) - height(old) 88 else: 89 # we know that for insertion we don't increase in height 90 hdiff = 0 91 return hdiff, new 92 elif lb == 0: 93 # can only happen in deletion 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: # lb == +1 99 # lrv will be the new root 100 lrl, lrv, lrr, lrb = lr 101 if lrb == 0: # lr is the new node 102 newleftb = newrightb = 0 103 elif lrb == -1: 104 newleftb = 0 105 newrightb = +1 106 else: # lrb == +1 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 # deletion; maybe we decreased in height 113 old = node(l, v, r, b) 114 hdiff += height(new) - height(old) 115 else: 116 # we know that for insertion we don't increase in height 117 hdiff = 0 118 119 return hdiff, new 120 elif b > 1: 121 # rotate left 122 rl, rv, rr, rb = r 123 if rb == +1: 124 # easy case, rv is new root 125 new = node(node(l, v, rl, 0), rv, rr, 0) 126 if hdiff <= 0: 127 # deletion; maybe we decreased in height 128 old = node(l, v, r, b) 129 hdiff += height(new) - height(old) 130 else: 131 # we know that for insertion we don't increase in height 132 hdiff = 0 133 return hdiff, new 134 elif rb == 0: 135 # can only happen in deletion 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: # rb == -1 141 # rlv will be the new root 142 rll, rlv, rlr, rlb = rl 143 if rlb == 0: # rl is the new node 144 newleftb = newrightb = 0 145 elif rlb == +1: 146 newleftb = -1 147 newrightb = 0 148 else: # rlb == -1 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 # deletion; maybe we decreased in height 155 old = node(l, v, r, b) 156 hdiff += height(new) - height(old) 157 else: 158 # we know that for insertion we don't increase in height 159 hdiff = 0 160 return hdiff, new 161 else: 162 return hdiff, node(l, v, r, b)
163 164
165 -def insert(tree, value):
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
190 -def delete(tree, value):
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 # overall height only changes if left was taller before 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 # overall height only changes if left was 217 # taller before 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 # overall height only changes if right was 226 # taller before 227 hdiff = 0 228 b -= 1 229 return _balance(hdiff, l, v, newr, b) 230 else: 231 # we have found the node! 232 if r is None: 233 # no right link, just replace with left 234 return -1, l 235 else: 236 newv, hdiff, newr = popmin(r) 237 if hdiff != 0: 238 if b <= 0: 239 # overall height only changes if right was 240 # taller before 241 hdiff = 0 242 b -= 1 243 return _balance(hdiff, l, newv, newr, b) 244 245
246 -def lookup(tree, value):
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
261 -def iterate(tree):
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
273 -def iteratereversed(tree):
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
285 -class AVLTree(object):
286
287 - def __init__(self, seq=()):
288 self._len = len(seq) 289 self.tree = fromseq(seq)
290
291 - def insert(self, value):
292 _, self.tree = insert(self.tree, value) 293 self._len += 1
294
295 - def delete(self, value):
296 _, self.tree = delete(self.tree, value) 297 self._len -= 1
298
299 - def __contains__(self, value):
300 return bool(lookup(self.tree, value))
301
302 - def __len__(self):
303 return self._len
304
305 - def __iter__(self):
306 return iterate(self.tree)
307
308 - def iterreversed(self):
309 return iteratereversed(self.tree)
310