third_party/bintrees/bintrees/avltree.py

   1 #!/usr/bin/env python
   2 #coding:utf-8
   3 # Author:  mozman (python version)
   4 # Purpose: avl tree module (Julienne Walker's unbounded none recursive algorithm)
   5 # source: http://eternallyconfuzzled.com/tuts/datastructures/jsw_tut_avl.aspx
   6 # Created: 01.05.2010
   7 # Copyright (c) 2010-2013 by Manfred Moitzi
   8 # License: MIT License
   9
  10 # Conclusion of Julienne Walker
  11
  12 # AVL trees are about as close to optimal as balanced binary search trees can
  13 # get without eating up resources. You can rest assured that the O(log N)
  14 # performance of binary search trees is guaranteed with AVL trees, but the extra
  15 # bookkeeping required to maintain an AVL tree can be prohibitive, especially
  16 # if deletions are common. Insertion into an AVL tree only requires one single
  17 # or double rotation, but deletion could perform up to O(log N) rotations, as
  18 # in the example of a worst case AVL (ie. Fibonacci) tree. However, those cases
  19 # are rare, and still very fast.
  20
  21 # AVL trees are best used when degenerate sequences are common, and there is
  22 # little or no locality of reference in nodes. That basically means that
  23 # searches are fairly random. If degenerate sequences are not common, but still
  24 # possible, and searches are random then a less rigid balanced tree such as red
  25 # black trees or Andersson trees are a better solution. If there is a significant
  26 # amount of locality to searches, such as a small cluster of commonly searched
  27 # items, a splay tree is theoretically better than all of the balanced trees
  28 # because of its move-to-front design.
  29
  30 from __future__ import absolute_import
  31
  32 from .treemixin import TreeMixin
  33 from array import array
  34
  35 __all__ = ['AVLTree']
  36
  37 MAXSTACK = 32
  38
  39
  40 class Node(object):
  41     """ Internal object, represents a treenode """
  42     __slots__ = ['left', 'right', 'balance', 'key', 'value']
  43
  44     def __init__(self, key=None, value=None):
  45         self.left = None
  46         self.right = None
  47         self.key = key
  48         self.value = value
  49         self.balance = 0
  50
  51     def __getitem__(self, key):
  52         """ x.__getitem__(key) <==> x[key], where key is 0 (left) or 1 (right) """
  53         return self.left if key == 0 else self.right
  54
  55     def __setitem__(self, key, value):
  56         """ x.__setitem__(key, value) <==> x[key]=value, where key is 0 (left) or 1 (right) """
  57         if key == 0:
  58             self.left = value
  59         else:
  60             self.right = value
  61
  62     def free(self):
  63         """Remove all references."""
  64         self.left = None
  65         self.right = None
  66         self.key = None
  67         self.value = None
  68
  69
  70 def height(node):
  71     return node.balance if node is not None else -1
  72
  73
  74 def jsw_single(root, direction):
  75     other_side = 1 - direction
  76     save = root[other_side]
  77     root[other_side] = save[direction]
  78     save[direction] = root
  79     rlh = height(root.left)
  80     rrh = height(root.right)
  81     slh = height(save[other_side])
  82     root.balance = max(rlh, rrh) + 1
  83     save.balance = max(slh, root.balance) + 1
  84     return save
  85
  86
  87 def jsw_double(root, direction):
  88     other_side = 1 - direction
  89     root[other_side] = jsw_single(root[other_side], other_side)
  90     return jsw_single(root, direction)
  91
  92
  93 class AVLTree(TreeMixin):
  94     """
  95     AVLTree implements a balanced binary tree with a dict-like interface.
  96
  97     see: http://en.wikipedia.org/wiki/AVL_tree
  98
  99     In computer science, an AVL tree is a self-balancing binary search tree, and
 100     it is the first such data structure to be invented. In an AVL tree, the
 101     heights of the two child subtrees of any node differ by at most one;
 102     therefore, it is also said to be height-balanced. Lookup, insertion, and
 103     deletion all take O(log n) time in both the average and worst cases, where n
 104     is the number of nodes in the tree prior to the operation. Insertions and
 105     deletions may require the tree to be rebalanced by one or more tree rotations.
 106
 107     The AVL tree is named after its two inventors, G.M. Adelson-Velskii and E.M.
 108     Landis, who published it in their 1962 paper "An algorithm for the
 109     organization of information."
 110
 111     AVLTree() -> new empty tree.
 112     AVLTree(mapping) -> new tree initialized from a mapping
 113     AVLTree(seq) -> new tree initialized from seq [(k1, v1), (k2, v2), ... (kn, vn)]
 114
 115     see also TreeMixin() class.
 116
 117     """
 118     def __init__(self, items=None):
 119         """ x.__init__(...) initializes x; see x.__class__.__doc__ for signature """
 120         self._root = None
 121         self._count = 0
 122         if items is not None:
 123             self.update(items)
 124
 125     def clear(self):
 126         """ T.clear() -> None.  Remove all items from T. """
 127         def _clear(node):
 128             if node is not None:
 129                 _clear(node.left)
 130                 _clear(node.right)
 131                 node.free()
 132         _clear(self._root)
 133         self._count = 0
 134         self._root = None
 135
 136     @property
 137     def count(self):
 138         """ count of items """
 139         return self._count
 140
 141     @property
 142     def root(self):
 143         """ root node of T """
 144         return self._root
 145
 146     def _new_node(self, key, value):
 147         """ Create a new treenode """
 148         self._count += 1
 149         return Node(key, value)
 150
 151     def insert(self, key, value):
 152         """ T.insert(key, value) <==> T[key] = value, insert key, value into Tree """
 153         if self._root is None:
 154             self._root = self._new_node(key, value)
 155         else:
 156             node_stack = []  # node stack
 157             dir_stack = array('I')  # direction stack
 158             done = False
 159             top = 0
 160             node = self._root
 161             # search for an empty link, save path
 162             while True:
 163                 if key == node.key:  # update existing item
 164                     node.value = value
 165                     return
 166                 direction = 1 if key > node.key else 0
 167                 dir_stack.append(direction)
 168                 node_stack.append(node)
 169                 if node[direction] is None:
 170                     break
 171                 node = node[direction]
 172
 173             # Insert a new node at the bottom of the tree
 174             node[direction] = self._new_node(key, value)
 175
 176             # Walk back up the search path
 177             top = len(node_stack) - 1
 178             while (top >= 0) and not done:
 179                 direction = dir_stack[top]
 180                 other_side = 1 - direction
 181                 topnode = node_stack[top]
 182                 left_height = height(topnode[direction])
 183                 right_height = height(topnode[other_side])
 184
 185                 # Terminate or rebalance as necessary */
 186                 if left_height - right_height == 0:
 187                     done = True
 188                 if left_height - right_height >= 2:
 189                     a = topnode[direction][direction]
 190                     b = topnode[direction][other_side]
 191
 192                     if height(a) >= height(b):
 193                         node_stack[top] = jsw_single(topnode, other_side)
 194                     else:
 195                         node_stack[top] = jsw_double(topnode, other_side)
 196
 197                     # Fix parent
 198                     if top != 0:
 199                         node_stack[top - 1][dir_stack[top - 1]] = node_stack[top]
 200                     else:
 201                         self._root = node_stack[0]
 202                     done = True
 203
 204                 # Update balance factors
 205                 topnode = node_stack[top]
 206                 left_height = height(topnode[direction])
 207                 right_height = height(topnode[other_side])
 208
 209                 topnode.balance = max(left_height, right_height) + 1
 210                 top -= 1
 211
 212     def remove(self, key):
 213         """ T.remove(key) <==> del T[key], remove item <key> from tree """
 214         if self._root is None:
 215             raise KeyError(str(key))
 216         else:
 217             node_stack = [None] * MAXSTACK  # node stack
 218             dir_stack = array('I', [0] * MAXSTACK)  # direction stack
 219             top = 0
 220             node = self._root
 221
 222             while True:
 223                 # Terminate if not found
 224                 if node is None:
 225                     raise KeyError(str(key))
 226                 elif node.key == key:
 227                     break
 228
 229                 # Push direction and node onto stack
 230                 direction = 1 if key > node.key else 0
 231                 dir_stack[top] = direction
 232
 233                 node_stack[top] = node
 234                 node = node[direction]
 235                 top += 1
 236
 237             # Remove the node
 238             if (node.left is None) or (node.right is None):
 239                 # Which child is not null?
 240                 direction = 1 if node.left is None else 0
 241
 242                 # Fix parent
 243                 if top != 0:
 244                     node_stack[top - 1][dir_stack[top - 1]] = node[direction]
 245                 else:
 246                     self._root = node[direction]
 247                 node.free()
 248                 self._count -= 1
 249             else:
 250                 # Find the inorder successor
 251                 heir = node.right
 252
 253                 # Save the path
 254                 dir_stack[top] = 1
 255                 node_stack[top] = node
 256                 top += 1
 257
 258                 while heir.left is not None:
 259                     dir_stack[top] = 0
 260                     node_stack[top] = heir
 261                     top += 1
 262                     heir = heir.left
 263
 264                 # Swap data
 265                 node.key = heir.key
 266                 node.value = heir.value
 267
 268                 # Unlink successor and fix parent
 269                 xdir = 1 if node_stack[top - 1].key == node.key else 0
 270                 node_stack[top - 1][xdir] = heir.right
 271                 heir.free()
 272                 self._count -= 1
 273
 274             # Walk back up the search path
 275             top -= 1
 276             while top >= 0:
 277                 direction = dir_stack[top]
 278                 other_side = 1 - direction
 279                 topnode = node_stack[top]
 280                 left_height = height(topnode[direction])
 281                 right_height = height(topnode[other_side])
 282                 b_max = max(left_height, right_height)
 283
 284                 # Update balance factors
 285                 topnode.balance = b_max + 1
 286
 287                 # Terminate or rebalance as necessary
 288                 if (left_height - right_height) == -1:
 289                     break
 290                 if (left_height - right_height) <= -2:
 291                     a = topnode[other_side][direction]
 292                     b = topnode[other_side][other_side]
 293                     if height(a) <= height(b):
 294                         node_stack[top] = jsw_single(topnode, direction)
 295                     else:
 296                         node_stack[top] = jsw_double(topnode, direction)
 297                     # Fix parent
 298                     if top != 0:
 299                         node_stack[top - 1][dir_stack[top - 1]] = node_stack[top]
 300                     else:
 301                         self._root = node_stack[0]
 302                 top -= 1