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