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  """ 
 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
35 -def height(tree):
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
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:'--',-1:'-',0:'0',1:'+',2:'++'}.get(b,'?') 54 print '%s%s: %r' % (' '*level, bchr, v) 55 debug(r, level+1)
56
57 -def fromseq(seq):
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
64 -def _balance(hdiff, l, v, r, b):
65 """Internal method to rebalance an AVL tree, called as needed.""" 66 # if we have to rebalance, in the end the node has balance 0; 67 # for details see GNU libavl docs 68 if b < -1: 69 # rotate right 70 ll, lv, lr, lb = l 71 if lb == -1: 72 # easy case, lv is new root 73 new = node(ll, lv, node(lr, v, r, 0), 0) 74 if hdiff <= 0: 75 # deletion; maybe we decreased in height 76 old = node(l, v, r, b) 77 hdiff += height(new) - height(old) 78 else: 79 # we know that for insertion we don't increase in height 80 hdiff = 0 81 return hdiff, new 82 elif lb == 0: 83 # can only happen in deletion 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: # lb == +1 89 # lrv will be the new root 90 lrl, lrv, lrr, lrb = lr 91 if lrb == 0: # lr is the new node 92 newleftb = newrightb = 0 93 elif lrb == -1: 94 newleftb = 0 95 newrightb = +1 96 else: # lrb == +1 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 # deletion; maybe we decreased in height 103 old = node(l, v, r, b) 104 hdiff += height(new) - height(old) 105 else: 106 # we know that for insertion we don't increase in height 107 hdiff = 0 108 109 return hdiff, new 110 elif b > 1: 111 # rotate left 112 rl, rv, rr, rb = r 113 if rb == +1: 114 # easy case, rv is new root 115 new = node(node(l, v, rl, 0), rv, rr, 0) 116 if hdiff <= 0: 117 # deletion; maybe we decreased in height 118 old = node(l, v, r, b) 119 hdiff += height(new) - height(old) 120 else: 121 # we know that for insertion we don't increase in height 122 hdiff = 0 123 return hdiff, new 124 elif rb == 0: 125 # can only happen in deletion 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: # rb == -1 131 # rlv will be the new root 132 rll, rlv, rlr, rlb = rl 133 if rlb == 0: # rl is the new node 134 newleftb = newrightb = 0 135 elif rlb == +1: 136 newleftb = -1 137 newrightb = 0 138 else: # rlb == -1 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 # deletion; maybe we decreased in height 145 old = node(l, v, r, b) 146 hdiff += height(new) - height(old) 147 else: 148 # we know that for insertion we don't increase in height 149 hdiff = 0 150 return hdiff, new 151 else: 152 return hdiff, node(l, v, r, b)
153
154 -def insert(tree, value):
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
178 -def delete(tree, value):
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 # overall height only changes if left was taller before 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 # overall height only changes if left was 204 # taller before 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 # overall height only changes if right was 213 # taller before 214 hdiff = 0 215 b -= 1 216 return _balance(hdiff, l, v, newr, b) 217 else: 218 # we have found the node! 219 if r is None: 220 # no right link, just replace with left 221 return -1, l 222 else: 223 newv, hdiff, newr = popmin(r) 224 if hdiff != 0: 225 if b <= 0: 226 # overall height only changes if right was 227 # taller before 228 hdiff = 0 229 b -= 1 230 return _balance(hdiff, l, newv, newr, b) 231
232 -def lookup(tree, value):
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
246 -def iterate(tree):
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
257 -def iteratereversed(tree):
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
268 -class AVLTree(object):
269 - def __init__(self, seq=()):
270 self._len = len(seq) 271 self.tree = fromseq(seq)
272
273 - def insert(self, value):
274 _, self.tree = insert(self.tree, value) 275 self._len += 1
276
277 - def delete(self, value):
278 _, self.tree = delete(self.tree, value) 279 self._len -= 1
280
281 - def __contains__(self, value):
282 return bool(lookup(self.tree, value))
283
284 - def __len__(self):
285 return self._len
286
287 - def __iter__(self):
288 return iterate(self.tree)
289
290 - def iterreversed(self):
291 return iteratereversed(self.tree)
292