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 30from __future__ import absolute_import 31 32from .treemixin import TreeMixin 33from array import array 34 35__all__ = ['AVLTree'] 36 37MAXSTACK = 32 38 39 40class 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 70def height(node): 71 return node.balance if node is not None else -1 72 73 74def 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 87def 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 93class 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 303