algorithm/avltree.py at master · pihou/algorithm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
#!-*- coding=utf8 -*-
from graphviz import Digraph

import queue
import random

class Node(object):
    def __init__(self, value):
        self.left = None
        self.right = None
        self.factor = 0
        self.parent = None
        self.value = value

    def get_parent(self):
        return self.parent

    def get_sibling(self):
        parent = self.parent
        if not parent:
            return None
        if parent.left is self:
            return parent.right
        return parent.left

    def get_grandparent(self):
        parent = self.parent
        if not parent:
            return None
        return parent.parent

    def get_uncle(self):
        parent = self.parent
        if not parent:
            return None
        if not parent.parent:
            return None
        if parent is parent.parent.left:
            return parent.parent.right
        return parent.parent.left

    def __str__(self):
        p = self.parent.value if self.parent else None
        return "parent: %s self:%s factor:%s"%(p, self.value, self.factor)

    def __repr__(self):
        return str(self)


class AVLTree(object):
    def __init__(self):
        self.root = None

    def rotate_left(self, n):
        p = n.get_parent()
        r = n.right
        assert(not r is None)

        n.right = r.left
        if r.left:
            r.left.parent = n
        n.parent = r
        r.left = n
        r.parent = p
        if not p:
            self.root = r
        elif p.left is n:
            p.left = r
        else:
            p.right = r

    def rotate_right(self, n):
        p = n.get_parent()
        l = n.left
        assert(not l is None)

        n.left = l.right
        if l.right:
            l.right.parent = n
        n.parent = l
        l.right = n
        l.parent = p
        if not p:
            self.root = l
        elif p.left is n:
            p.left = l
        else:
            p.right = l

#---------------------------------------------------------
    def _push_node(self, node):
        n = self.root
        if not n:
            self.root = node
            return
        while True:
            if node.value <= n.value:
                if n.left is None:
                    n.left = node
                    node.parent = n
                    break
                n = n.left
            else:
                if n.right is None:
                    n.right = node
                    node.parent = n
                    break
                n = n.right

    def _balance_tree(self, node):
        n = node
        if n.get_sibling():
            n.parent.factor = 0
            return
        while n.parent:
            p = n.parent
            if p.left is n:
                p.factor -= 1
            else:
                p.factor += 1
            if not p.factor:
                break
            if abs(p.factor) <= 1:
                n = p
                continue
            if p.factor < -1 and n.factor < 0:
                n.factor = 0
                p.factor = 0
                self.rotate_right(p)
                break
            if p.factor > 1 and n.factor > 0:
                n.factor = 0
                p.factor = 0
                self.rotate_left(p)
                break
            if p.factor < -1 and n.factor > 0:
                r = n.right
                factor = r.factor
                n.factor = 0 if factor <= 0 else -1
                self.rotate_left(n)
                r.factor = 0
                p.factor = 0 if factor >= 0 else 1
                self.rotate_right(p)
                break
            if p.factor > 1 and n.factor < 0:
                l = n.left
                factor = l.factor
                n.factor = 0 if factor >= 0 else 1
                self.rotate_right(n)
                l.factor = 0
                p.factor = 0 if factor <= 0 else -1
                self.rotate_left(p)
                break
            raise Erro

    def add(self, node):
        self._push_node(node)
        self._balance_tree(node)


if __name__ == "__main__":
    def print_tree(tree, dot, index):
        def get_sign(index, node):
            return str(index) + str(id(node))
        q = queue.Queue()
        q.put(tree.root)
        while not q.empty():
            node = q.get()
            sign = get_sign(index, node)
            dot.node(sign, "%s(%s)"%(node.value, node.factor))
            if node.left:
                dot.edge(sign, get_sign(index, node.left))
                q.put(node.left)
            if node.right:
                dot.edge(sign, get_sign(index, node.right))
                q.put(node.right)
        return dot

    dots = Digraph()
    tree = AVLTree()
    for i, value in enumerate(xrange(24)):
        value = random.randint(0, 10000)
        tree.add(Node(value))
        dot = Digraph()
        print_tree(tree, dot, i)
        dots.subgraph(graph=dot)
    dots.render(view=True)