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: 6964 $" 
 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 -def height(tree):
38 """Return the height of an AVL tree. Relies on the balance factors 39 being consistent.""" 40 if tree is None: 41 return 0 42 else: 43 l, v, r, b = tree 44 if b <= 0: 45 return height(l) + 1 46 else: 47 return height(r) + 1
48
49 -def debug(tree, level=0):
50 """Print out a debugging representation of an AVL tree.""" 51 if tree is None: 52 return 53 l, v, r, b = tree 54 debug(l, level+1) 55 bchr = {-2:'--',-1:'-',0:'0',1:'+',2:'++'}.get(b,'?') 56 print '%s%s: %r' % (' '*level, bchr, v) 57 debug(r, level+1)
58
59 -def fromseq(seq):
60 """Populate and return an AVL tree from an iterable sequence.""" 61 t = None 62 for x in seq: 63 _, t = insert(t, x) 64 return t
65
66 -def _balance(hdiff, l, v, r, b):
67 """Internal method to rebalance an AVL tree, called as needed.""" 68 # if we have to rebalance, in the end the node has balance 0; 69 # for details see GNU libavl docs 70 if b < -1: 71 # rotate right 72 ll, lv, lr, lb = l 73 if lb == -1: 74 # easy case, lv is new root 75 new = node(ll, lv, node(lr, v, r, 0), 0) 76 if hdiff <= 0: 77 # deletion; maybe we decreased in height 78 old = node(l, v, r, b) 79 hdiff += height(new) - height(old) 80 else: 81 # we know that for insertion we don't increase in height 82 hdiff = 0 83 return hdiff, new 84 elif lb == 0: 85 # can only happen in deletion 86 new = node(ll, lv, node(lr, v, r, -1), +1) 87 old = node(l, v, r, b) 88 hdiff += height(new) - height(old) 89 return hdiff, new 90 else: # lb == +1 91 # lrv will be the new root 92 lrl, lrv, lrr, lrb = lr 93 if lrb == 0: # lr is the new node 94 newleftb = newrightb = 0 95 elif lrb == -1: 96 newleftb = 0 97 newrightb = +1 98 else: # lrb == +1 99 newleftb = -1 100 newrightb = 0 101 new = node(node(ll, lv, lrl, newleftb), lrv, 102 node(lrr, v, r, newrightb), 0) 103 if hdiff <= 0: 104 # deletion; maybe we decreased in height 105 old = node(l, v, r, b) 106 hdiff += height(new) - height(old) 107 else: 108 # we know that for insertion we don't increase in height 109 hdiff = 0 110 111 return hdiff, new 112 elif b > 1: 113 # rotate left 114 rl, rv, rr, rb = r 115 if rb == +1: 116 # easy case, rv is new root 117 new = node(node(l, v, rl, 0), rv, rr, 0) 118 if hdiff <= 0: 119 # deletion; maybe we decreased in height 120 old = node(l, v, r, b) 121 hdiff += height(new) - height(old) 122 else: 123 # we know that for insertion we don't increase in height 124 hdiff = 0 125 return hdiff, new 126 elif rb == 0: 127 # can only happen in deletion 128 new = node(node(l, v, rl, +1), rv, rr, -1) 129 old = node(l, v, r, b) 130 hdiff += height(new) - height(old) 131 return hdiff, new 132 else: # rb == -1 133 # rlv will be the new root 134 rll, rlv, rlr, rlb = rl 135 if rlb == 0: # rl is the new node 136 newleftb = newrightb = 0 137 elif rlb == +1: 138 newleftb = -1 139 newrightb = 0 140 else: # rlb == -1 141 newleftb = 0 142 newrightb = +1 143 new = node(node(l, v, rll, newleftb), rlv, 144 node(rlr, rv, rr, newrightb), 0) 145 if hdiff <= 0: 146 # deletion; maybe we decreased in height 147 old = node(l, v, r, b) 148 hdiff += height(new) - height(old) 149 else: 150 # we know that for insertion we don't increase in height 151 hdiff = 0 152 return hdiff, new 153 else: 154 return hdiff, node(l, v, r, b)
155
156 -def insert(tree, value):
157 """Insert a value into an AVL tree. Returns a tuple of 158 (heightdifference, tree). The original tree is unmodified.""" 159 if tree is None: 160 return 1, (None, value, None, 0) 161 else: 162 l, v, r, b = tree 163 if value < v: 164 hdiff, newl = insert(l, value) 165 if hdiff > 0: 166 if b > 0: 167 hdiff = 0 168 b -= 1 169 return _balance(hdiff, newl, v, r, b) 170 elif value > v: 171 hdiff, newr = insert(r, value) 172 if hdiff > 0: 173 if b < 0: 174 hdiff = 0 175 b += 1 176 return _balance(hdiff, l, v, newr, b) 177 else: 178 raise ValueError('tree already has value %r' % (value,))
179
180 -def delete(tree, value):
181 """Delete a value from an AVL tree. Like L{insert}, returns a tuple 182 of (heightdifference, tree). The original tree is unmodified.""" 183 def popmin((l, v, r, b)): 184 if l is None: 185 minv = v 186 return minv, -1, r 187 else: 188 minv, hdiff, newl = popmin(l) 189 if hdiff != 0: 190 if b >= 0: 191 # overall height only changes if left was taller before 192 hdiff = 0 193 b += 1 194 195 return (minv,) + _balance(hdiff, newl, v, r, b)
196 197 if tree is None: 198 raise ValueError('tree has no value %r' % (value,)) 199 else: 200 l, v, r, b = tree 201 if value < v: 202 hdiff, newl = delete(l, value) 203 if hdiff != 0: 204 if b >= 0: 205 # overall height only changes if left was 206 # taller before 207 hdiff = 0 208 b += 1 209 return _balance(hdiff, newl, v, r, b) 210 elif value > v: 211 hdiff, newr = delete(r, value) 212 if hdiff != 0: 213 if b <= 0: 214 # overall height only changes if right was 215 # taller before 216 hdiff = 0 217 b -= 1 218 return _balance(hdiff, l, v, newr, b) 219 else: 220 # we have found the node! 221 if r is None: 222 # no right link, just replace with left 223 return -1, l 224 else: 225 newv, hdiff, newr = popmin(r) 226 if hdiff != 0: 227 if b <= 0: 228 # overall height only changes if right was 229 # taller before 230 hdiff = 0 231 b -= 1 232 return _balance(hdiff, l, newv, newr, b) 233
234 -def lookup(tree, value):
235 """Look up a node in an AVL tree. Returns a node tuple or False if 236 the value was not found.""" 237 if tree is None: 238 return False 239 else: 240 l, v, r, b = tree 241 if value < v: 242 return lookup(l, v) 243 elif value > v: 244 return lookup(r, v) 245 else: 246 return tree
247
248 -def iterate(tree):
249 """Iterate over an AVL tree, starting with the lowest-ordered 250 value.""" 251 if tree is not None: 252 l, v, r, b = tree 253 for x in iterate(l): 254 yield x 255 yield v 256 for x in iterate(r): 257 yield x
258
259 -def iteratereversed(tree):
260 """Iterate over an AVL tree, starting with the highest-ordered 261 value.""" 262 if tree is not None: 263 l, v, r, b = tree 264 for x in iteratereversed(r): 265 yield x 266 yield v 267 for x in iteratereversed(l): 268 yield x
269
270 -class AVLTree(object):
271 - def __init__(self, seq=()):
272 self._len = len(seq) 273 self.tree = fromseq(seq)
274
275 - def insert(self, value):
276 _, self.tree = insert(self.tree, value) 277 self._len += 1
278
279 - def delete(self, value):
280 _, self.tree = delete(self.tree, value) 281 self._len -= 1
282
283 - def __contains__(self, value):
284 return bool(lookup(self.tree, value))
285
286 - def __len__(self):
287 return self._len
288
289 - def __iter__(self):
290 return iterate(self.tree)
291
292 - def iterreversed(self):
293 return iteratereversed(self.tree)
294