CS 공부 - [AVL 트리]

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
# AVL트리 안에 노드를 삽입하는 코드

# 일반적인 트리 노드 클래스
class TreeNode(object):
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
        self.height = 1

# 삽입 연산을 지원하는 AVL 트리 클래스
class AVL_Tree(object):
    # 서브트리에 키를 삽입하고
    # 서브트리의 새로운 루트를 리턴하는 재귀함수 입니다.
    # 과정은 다음과 같습니다.
    # 1. 노드를 삽입하고, 루트로 올라가면서 높이를 업데이트 합니다.
    # 2. 업데이트하려는 노드의 양쪽 자식의 높이 차이가 2가 된다면,
    #    이 노드를 X라 하고 X의 두 자식 중 높이가 큰 쪽을 y, y의 두 자식 중 높이가 큰 쪽을 z라고 합니다.
    # 3. x, y, z를 대소관계에 따라 잘 재배치합니다.
    def insert(self, root, key):

        # Step 1 - 일반적인 이진탐색트리 만들기
        if not root:
            return TreeNode(key)
        elif key < root.val:
            root.left = self.insert(root.left, key)
        else:
            root.right = self.insert(root.right, key)

        # Step 2 - 조상노드의 높이 업데이트
        root.height = 1 + max(self.getHeight(root.left),
                           self.getHeight(root.right))

        # Step 3 - 밸런스 요소 찾기
        balance = self.getBalance(root)

        # Step 4 - 노드가 불균형이면 4가지 케이스중 하나를 실행
        # Case 1 - Left Left
        if balance > 1 and key < root.left.val:
            return self.rightRotate(root)

        # Case 2 - Right Right
        if balance < -1 and key > root.right.val:
            return self.leftRotate(root)

        # Case 3 - Left Right
        if balance > 1 and key > root.left.val:
            root.left = self.leftRotate(root.left)
            return self.rightRotate(root)

        # Case 4 - Right Left
        if balance < -1 and key < root.right.val:
            root.right = self.rightRotate(root.right)
            return self.leftRotate(root)

        return root

    def leftRotate(self, z):

        y = z.right
        T2 = y.left

        # 회전 수행
        y.left = z
        z.right = T2

        # 높이 업데이트
        z.height = 1 + max(self.getHeight(z.left),
                         self.getHeight(z.right))
        y.height = 1 + max(self.getHeight(y.left),
                         self.getHeight(y.right))

        # 새로운 루트 리턴
        return y

    def rightRotate(self, z):

        y = z.left
        T3 = y.right

        # 회전 수행
        y.right = z
        z.left = T3

        # 높이 업데이트
        z.height = 1 + max(self.getHeight(z.left),
                        self.getHeight(z.right))
        y.height = 1 + max(self.getHeight(y.left),
                        self.getHeight(y.right))

        # 새로운 루트 리턴
        return y

    def getHeight(self, root):
        if not root:
            return 0

        return root.height

    def getBalance(self, root):
        if not root:
            return 0

        return self.getHeight(root.left) - self.getHeight(root.right)

    def preOrder(self, root):

        if not root:
            return

        print("{0} ".format(root.val), end="")
        self.preOrder(root.left)
        self.preOrder(root.right)


myTree = AVL_Tree()
root = None

root = myTree.insert(root, 10)
root = myTree.insert(root, 20)
root = myTree.insert(root, 30)
root = myTree.insert(root, 40)
root = myTree.insert(root, 50)
root = myTree.insert(root, 25)

"""AVL 트리는 최종적으로 다음과 같이 보일것입니다.
            30
           /  \
         20   40
        /  \     \
       10  25    50"""

# 전위 순회 출력
myTree.preOrder(root)

# This code is contributed by Ajitesh Pathak