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