[Python] 이진 탐색 트리(Binary Search Tree)

 

이진 트리

이진 탐색 트리에 들어가기에 앞서 먼저 이진 트리에 대해 간단하게 알아보자. 이진 트리란, 트리의 종류로 각 노드가 최대 2개의 자식 노드를 갖는 트리를 말한다. 즉, 이진 트리의 각 노드는 자식 노드를 갖고 있지 않을수도 있으며, 1개를 가질 수도, 2개를 가질 수도 있다.

 

이진 탐색 트리

이진 탐색 트리는 이진 트리에서 어떠한 규칙에 따라 나열한 트리이다. 이 규칙은 모든 노드에 대해서 왼쪽 노드보다 오른쪽 노드가 더 크게 나열하는 것이다. 아래 그림을 보자. 

 

 

21이 루트 노드이고, 다음으로 28을 트리에 삽입을 할 때 21보다 큰 값이므로 오른쪽으로 간다. 그 다음 값인 14는 21보다 작으므로 왼쪽으로 간다. 그리고 아래 그림이 이진 탐색 트리에서 값을 찾는 것을 애니메이션으로 나타낸 것이다. 이진 탐색 트리는 평균 O(logN)의 시간복잡도를 가지며 트리가 한쪽으로만 치우친 최악의 경우 O(N)의 시간복잡도를 가진다. 

 

 

파이썬 구현

파이썬으로 이진탐색트리를 구현하는 데에는 재귀적으로 구현하는 방법과 단순 구현하는 방법이 있다. 재귀적으로 구현하는 방법이 코드도 훨씬 간결하고 많이 쓰이는 것 같다. 이 글에서는 두 가지 방법 다 적어보려 한다. 먼저 단순 구현을 한 코드이다.

 

# 노드 생성
class Node:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

 

먼저 위와 같이 루트 노드를 생성한다. 현재는 루트 노드에 자식 노드가 없으므로 left와 right에 None을 넣어준다. 그리고 insert를 통해 하나씩 노드를 삽입해보자.

 

# 노드 삽입
class BST:
    def __init__(self, root):
        self.root = root

    def insert(self, value):
        self.current_node = self.root
        while True:
            if value < self.current_node.value:
                if self.current_node.left != None:
                    self.current_node = self.current_node.left
                else:
                    self.current_node.left = Node(value)
                    break
            else:
                if self.current_node.right != None:
                    self.current_node = self.current_node.right
                else:
                    self.current_node.right = Node(value)
                    break

 

먼저, 아래와 같이 current_node를 루트 노드로 초기화 시킨다.

 

def insert(self, value):
        self.current_node = self.root

 

삽입하고자 하는 값이 current_node의 값보다 작을 때. 현재 노드의 왼쪽 자식 노드가 있는지 확인하고 있으면 current_node를 갱신한다. 없을 경우에는 current_node의 왼쪽 자식 노드에 삽입하고자 하는 노드를 삽입한다.

 

if value < self.current_node.value:
    if self.current_node.left != None:
        self.current_node = self.current_node.left
    else:
        self.current_node.left = Node(value)
        break

 

삽입하고자 하는 값이 current_node의 값보다 클 때. 작을 때의 경우의 반대로 구현한다.

 

else:
    if self.current_node.right != None:
        self.current_node = self.current_node.right
    else:
        self.current_node.right = Node(value)
        break

 

삽입을 한 후, 이진 탐색 트리에서 값을 탐색하는 search함수이다. 

 

# 노드 탐색
def search(self, value):
    self.current_node = self.root
    while self.current_node:
        if self.current_node.value == value:
            return True
        elif self.current_node.value > value:
            self.current_node = self.current_node.left
        else:
            self.current_node = self.current_node.right
    return False

 

루트 노드부터 탐색을 시작. current_node의 값이 찾고자 하는 값일 경우 True를 리턴한다. current_node의 값이 찾고자 하는 값보다 더 클 경우 찾고자 하는 값은 현재 노드의 왼쪽 노드에 있는 것이므로 current_node를 왼쪽 자식 노드인 current_node.left로 갱신해준다. 클 경우는 current_node.right로 갱신한다. 이렇게 while문을 돌면서 찾고자 하는 값을 찾으면 True를 리턴할 것이고, 못 찾을 경우 마지막 current_node에는 None값이 들어가므로 while문을 빠져나오게 된다. 그리고 False를 리턴한다. 그럼 이제 몇 가지 노드를 삽입해보고 노드가 잘 삽입이 되었는지 확인해보자.

 

root = Node(1)
bst = BST(root)
bst.insert(2)
bst.insert(7)
bst.insert(8)
bst.insert(10)
print(bst.search(2)) // True
print(bst.search(5)) // False
print(bst.search(7)) // True
print(bst.search(8)) // True
print(bst.search(15)) // False

 

다음으로 노드를 삭제하는 delete함수를 구현해야한다. 노드를 삭제하는 로직은 삽입이나 탐색보다 훨씬 복잡하다. 삭제할 노드가 자식노드를 하나도 갖고있지 않으면 괜찮지만, 자식노드를 갖고있고 그 자식노드가 또 다른 자식노드를 갖고있을 수 있기 때문이다. 그래서 구현을 할 때에도 그렇게 경우의 수를 나눠서 구현을 한다. 

 

def delete(self, value):
        # 삭제할 노드가 있는지 검사하고 없으면 False리턴
        # 검사를 한 후에는 삭제할 노드가 current_node, 삭제할 노드의 부모 노드가 parent가 된다.
        is_search = False
        self.current_node = self.root
        self.parent = self.root
        while self.current_node:
            if self.current_node.value == value:
                is_search = True
                break
            elif value < self.current_node.value:
                self.parent = self.current_node
                self.current_node = self.current_node.left
            else:
                self.parent = self.current_node
                self.current_node = self.current_node.right
        if is_search == False:
            return False
		
        # 삭제할 노드가 자식 노드를 갖고 있지 않을 때
        if self.current_node.left == None and self.current_node.right == None:
            if value < self.parent.value:
                self.parent.left = None
            else:
                self.parent.right = None
        
        # 삭제할 노드가 자식 노드를 한 개 가지고 있을 때(왼쪽 자식 노드)
        if self.current_node.left != None and self.current_node.right == None:
            if value < self.parent.value:
                self.parent.left = self.current_node.left
            else:
                self.parent.right = self.current_node.left
        
        # 삭제할 노드가 자식 노드를 한 개 가지고 있을 때(오른쪽 자식 노드)
        if self.current_node.left == None and self.current_node.right != None:
            if value < self.parent.value:
                self.parent.left = self.current_node.right
            else:
                self.parent.right = self.current_node.right                

        # 삭제할 노드가 자식 노드를 두 개 가지고 있을 때
        if self.current_node.left != None and self.current_node.right != None:
            self.change_node = self.current_node.right
            self.change_node_parent = self.current_node.right
            while self.change_node.left != None:
                self.change_node_parent = self.change_node
                self.change_node = self.change_node.left
            if self.change_node.right != None:
                self.change_node_parent.left = self.change_node.right
            else:
                self.change_node_parent.left = None
                
            if value < self.parent.value:
                self.parent.left = self.change_node
                self.change_node.right = self.current_node.right
                self.change_node.left = self.current_node.left
            else:
                self.parent.right = self.change_node
                self.change_node.left = self.current_node.left
                self.change_node.right = self.current_node.right

        return True

 

코드가 너무 복잡하다🤮그럼 이제 delete함수가 잘 구현되었는지 확인해보자. 

 

arr = [5, 2, 4, 22, 10, 12, 15, 60, 44, 9]
root = Node(30)
bst = BST(root)
for i in arr:
    bst.insert(i)
    
print(bst.search(22)) // True
print(bst.search(61)) // False
print(bst.search(60)) // True
print(bst.delete(60)) // True
print(bst.search(60)) // False
print(bst.delete(22)) // True
print(bst.delete(44)) // True
print(bst.search(22)) // False
print(bst.search(44)) // False

 

재귀를 이용한 코드

# 노드 생성과 삽입
class Node:
    def __init__(self,data):
        self.left = None
        self.right = None
        self.data = data
        
class BinarySearchTree(Node):
    def insert(self, data):
        self.root = self._insert_value(self.root, data)
        return self.root is not None

    def _insert_value(self, node, data):
        if node is None:
            node = Node(data)
        else:
            if data <= node.data:
                node.left = self._insert_value(node.left, data)
            else:
                node.right = self._insert_value(node.right, data)
        return node

 

# 노드 탐색
    def find(self, key):
        return self._find_value(self.root, key)

    def _find_value(self, root, key):
        if root is None or root.data == key:
            return root is not None
        elif key < root.data:
            return self._find_value(root.left, key)
        else:
            return self._find_value(root.right, key)

 

자식 노드가 2개일 경우, delete는 삭제할 노드와 오른쪽 서브 트리에서 가장 작은 노드의 값을 바꾸면 된다. 이 노드는 왼쪽 서브 트리에 있는 모든 값보다 크고 오른쪽 서브 트리에 있는 값 중 가장 작기 때문이다.

 

    def delete(self, key):
        self.root, deleted = self._delete_value(self.root, key)
        return deleted

    def _delete_value(self, node, key):
        if node is None:
            return node, False

        deleted = False
        # 해당 노드가 삭제할 노드일 경우
        if key == node.data:
            deleted = True
            # 삭제할 노드가 자식이 두개일 경우
            if node.left and node.right:
                # 오른쪽 서브 트리에서 가장 왼쪽에 있는 노드를 찾고 교체
                parent, child = node, node.right
                while child.left is not None:
                    parent, child = child, child.left
                child.left = node.left
                if parent != node:
                    parent.left = child.right
                    child.right = node.right
                node = child
            # 자식 노드가 하나일 경우 해당 노드와 교체
            elif node.left or node.right:
                node = node.left or node.right
            # 자식 노드가 없을 경우 그냥 삭제
            else:
                node = None
        elif key < node.data:
            node.left, deleted = self._delete_value(node.left, key)
        else:
            node.right, deleted = self._delete_value(node.right, key)
        return node, deleted

 

다음으로는 트리 순회 알고리즘에 대해 포스팅 할 예정이다🐳

 

 

참고

• https://austingwalters.com/binary-trees-traversals-everyday-algorithms/
• https://en.wikipedia.org/wiki/Binary_search_tree
• http://ejklike.github.io/2018/01/09/traversing-a-binary-tree-1.html