1#!/usr/bin/env python
2#coding:utf-8
3# Author:  mozman (python version)
4# Purpose: red-black tree module (Julienne Walker's none recursive algorithm)
5# source: http://eternallyconfuzzled.com/tuts/datastructures/jsw_tut_rbtree.aspx
6# Created: 01.05.2010
7# Copyright (c) 2010-2013 by Manfred Moitzi
8# License: MIT License
9
10# Conclusion of Julian Walker
11
12# Red black trees are interesting beasts. They're believed to be simpler than
13# AVL trees (their direct competitor), and at first glance this seems to be the
14# case because insertion is a breeze. However, when one begins to play with the
15# deletion algorithm, red black trees become very tricky. However, the
16# counterweight to this added complexity is that both insertion and deletion
17# can be implemented using a single pass, top-down algorithm. Such is not the
18# case with AVL trees, where only the insertion algorithm can be written top-down.
19# Deletion from an AVL tree requires a bottom-up algorithm.
20
21# So when do you use a red black tree? That's really your decision, but I've
22# found that red black trees are best suited to largely random data that has
23# occasional degenerate runs, and searches have no locality of reference. This
24# takes full advantage of the minimal work that red black trees perform to
25# maintain balance compared to AVL trees and still allows for speedy searches.
26
27# Red black trees are popular, as most data structures with a whimsical name.
28# For example, in Java and C++, the library map structures are typically
29# implemented with a red black tree. Red black trees are also comparable in
30# speed to AVL trees. While the balance is not quite as good, the work it takes
31# to maintain balance is usually better in a red black tree. There are a few
32# misconceptions floating around, but for the most part the hype about red black
33# trees is accurate.
34
35from __future__ import absolute_import
36
37from .treemixin import TreeMixin
38
39__all__ = ['RBTree']
40
41
42class Node(object):
43    """ Internal object, represents a treenode """
44    __slots__ = ['key', 'value', 'red', 'left', 'right']
45
46    def __init__(self, key=None, value=None):
47        self.key = key
48        self.value = value
49        self.red = True
50        self.left = None
51        self.right = None
52
53    def free(self):
54        self.left = None
55        self.right = None
56        self.key = None
57        self.value = None
58
59    def __getitem__(self, key):
60        """ x.__getitem__(key) <==> x[key], where key is 0 (left) or 1 (right) """
61        return self.left if key == 0 else self.right
62
63    def __setitem__(self, key, value):
64        """ x.__setitem__(key, value) <==> x[key]=value, where key is 0 (left) or 1 (right) """
65        if key == 0:
66            self.left = value
67        else:
68            self.right = value
69
70
71def is_red(node):
72    if (node is not None) and node.red:
73        return True
74    else:
75        return False
76
77
78def jsw_single(root, direction):
79    other_side = 1 - direction
80    save = root[other_side]
81    root[other_side] = save[direction]
82    save[direction] = root
83    root.red = True
84    save.red = False
85    return save
86
87
88def jsw_double(root, direction):
89    other_side = 1 - direction
90    root[other_side] = jsw_single(root[other_side], other_side)
91    return jsw_single(root, direction)
92
93
94class RBTree(TreeMixin):
95    """
96    RBTree implements a balanced binary tree with a dict-like interface.
97
98    see: http://en.wikipedia.org/wiki/Red_black_tree
99
100    A red-black tree is a type of self-balancing binary search tree, a data
101    structure used in computing science, typically used to implement associative
102    arrays. The original structure was invented in 1972 by Rudolf Bayer, who
103    called them "symmetric binary B-trees", but acquired its modern name in a
104    paper in 1978 by Leonidas J. Guibas and Robert Sedgewick. It is complex,
105    but has good worst-case running time for its operations and is efficient in
106    practice: it can search, insert, and delete in O(log n) time, where n is
107    total number of elements in the tree. Put very simply, a red-black tree is a
108    binary search tree which inserts and removes intelligently, to ensure the
109    tree is reasonably balanced.
110
111    RBTree() -> new empty tree.
112    RBTree(mapping) -> new tree initialized from a mapping
113    RBTree(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:  # Empty tree case
154            self._root = self._new_node(key, value)
155            self._root.red = False  # make root black
156            return
157
158        head = Node()  # False tree root
159        grand_parent = None
160        grand_grand_parent = head
161        parent = None  # parent
162        direction = 0
163        last = 0
164
165        # Set up helpers
166        grand_grand_parent.right = self._root
167        node = grand_grand_parent.right
168        # Search down the tree
169        while True:
170            if node is None:  # Insert new node at the bottom
171                node = self._new_node(key, value)
172                parent[direction] = node
173            elif is_red(node.left) and is_red(node.right):  # Color flip
174                node.red = True
175                node.left.red = False
176                node.right.red = False
177
178            # Fix red violation
179            if is_red(node) and is_red(parent):
180                direction2 = 1 if grand_grand_parent.right is grand_parent else 0
181                if node is parent[last]:
182                    grand_grand_parent[direction2] = jsw_single(grand_parent, 1 - last)
183                else:
184                    grand_grand_parent[direction2] = jsw_double(grand_parent, 1 - last)
185
186            # Stop if found
187            if key == node.key:
188                node.value = value  # set new value for key
189                break
190
191            last = direction
192            direction = 0 if key < node.key else 1
193            # Update helpers
194            if grand_parent is not None:
195                grand_grand_parent = grand_parent
196            grand_parent = parent
197            parent = node
198            node = node[direction]
199
200        self._root = head.right  # Update root
201        self._root.red = False  # make root black
202
203    def remove(self, key):
204        """ T.remove(key) <==> del T[key], remove item <key> from tree """
205        if self._root is None:
206            raise KeyError(str(key))
207        head = Node()  # False tree root
208        node = head
209        node.right = self._root
210        parent = None
211        grand_parent = None
212        found = None  # Found item
213        direction = 1
214
215        # Search and push a red down
216        while node[direction] is not None:
217            last = direction
218
219            # Update helpers
220            grand_parent = parent
221            parent = node
222            node = node[direction]
223
224            direction = 1 if key > node.key else 0
225
226            # Save found node
227            if key == node.key:
228                found = node
229
230            # Push the red node down
231            if not is_red(node) and not is_red(node[direction]):
232                if is_red(node[1 - direction]):
233                    parent[last] = jsw_single(node, direction)
234                    parent = parent[last]
235                elif not is_red(node[1 - direction]):
236                    sibling = parent[1 - last]
237                    if sibling is not None:
238                        if (not is_red(sibling[1 - last])) and (not is_red(sibling[last])):
239                            # Color flip
240                            parent.red = False
241                            sibling.red = True
242                            node.red = True
243                        else:
244                            direction2 = 1 if grand_parent.right is parent else 0
245                            if is_red(sibling[last]):
246                                grand_parent[direction2] = jsw_double(parent, last)
247                            elif is_red(sibling[1-last]):
248                                grand_parent[direction2] = jsw_single(parent, last)
249                            # Ensure correct coloring
250                            grand_parent[direction2].red = True
251                            node.red = True
252                            grand_parent[direction2].left.red = False
253                            grand_parent[direction2].right.red = False
254
255        # Replace and remove if found
256        if found is not None:
257            found.key = node.key
258            found.value = node.value
259            parent[int(parent.right is node)] = node[int(node.left is None)]
260            node.free()
261            self._count -= 1
262
263        # Update root and make it black
264        self._root = head.right
265        if self._root is not None:
266            self._root.red = False
267        if not found:
268            raise KeyError(str(key))