[Python 자료구조] 간단한 이진 탐색 트리(BST, Binary Search Tree) 구현하기

최근 자료구조와 알고리즘 강의를 듣고 있다.

이 강의를 듣고 나면 백준 알고리즘 문제도 조금씩 풀어볼 예정이다.

오늘 들은 강의 중, 이진 트리(BST)를 구현해보는 부분이 있었는데, 기본적인 내용이라 그런가 걱정했던 것과 달리 재밌었다.

아래의 코드는 강의에서 알려준 코드를 참고해 다시 나름대로 다시 작성해본 코드다.

기록 겸 남겨본다.

class Node:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

class NodeMgmt:
    def __init__(self, value):
        self.head = Node(value)
    
    # 노드 추가
    def insert(self, value):
        self.current_node = self.head
        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
                    
    # 탐색
    def search(self, value):
        self.searched = False
        self.parent_node = self.head
        self.current_node = self.head
        while self.current_node:
            if self.current_node.value == value:
                self.searched = True
                break
            else:
                self.parent_node = self.current_node
                if value < self.current_node.value:
                    self.current_node = self.current_node.left
                else:
                    self.current_node = self.current_node.right
    
    # 삭제
    def delete(self, value):
        # 삭제할 노드가 있는지부터 탐색
        self.search(value)

        if self.searched == False:
            return self.searched
        
        # case 구분 전, 삭제할 노드가 parent node의 왼쪽과 오른쪽이냐로 구분해준다
        if self.parent_node.value > value: # current node의 value가 parent node의 value보다 작은 경우. 즉, 왼쪽
            # case 1 : 삭제할 노드가 leaf node인 경우
            if self.current_node.left == None and self.current_node.right == None:
                self.parent_node.left == None
            # case 2 : 삭제할 노드에 child node가 하나 있는 경우
            elif self.current_node.left == None and self.current_node.right != None:
                self.parent_node.left = self.current_node.right
            elif self.current_node.left != None and self.current_node.right == None:
                self.parent_node.left = self.current_node.left
            # case 3 : 삭제할 노드에 child node가 두개 있는 경우
            else:
                self.child_parent_node = self.current_node
                self.child_node = self.current_node.right

                # 삭제할 노드의 오른쪽 자식 노드의 왼쪽 자식 중 가장 작은 값을 찾는다
                while self.child_node.left:
                    self.child_parent_node = self.child_node
                    self.child_node = self.child_node.left
                
                # 왼쪽 자식 노드 중 가장 작은 값을 가지는 노드에 오른쪽 자식이 있는 경우
                if self.child_node.right != None:
                    self.child_parent_node.left = self.child_node.right

                self.parent_node.left = self.child_node
                self.child_node.left = self.current_node.left
                self.child_node.right = self.current_node.right

        else: # current node의 value가 parent node의 value보다 크거나 같은 경우. 즉, 오른쪽
            # case 1 : 삭제할 노드가 leaf node인 경우
            if self.current_node.left == None and self.current_node.right == None:
                self.parent_node.right == None
            # case 2 : 삭제할 노드에 child node가 하나 있는 경우
            elif self.current_node.left == None and self.current_node.right != None:
                self.parent_node.right = self.current_node.right
            elif self.current_node.left != None and self.current_node.right == None:
                self.parent_node.right = self.current_node.left
            # case 3 : 삭제할 노드에 child node가 두개 있는 경우
            else:
                self.child_parent_node = self.current_node
                self.child_node = self.current_node.right

                # 삭제할 노드의 오른쪽 자식 노드의 왼쪽 자식 중 가장 작은 값을 찾는다
                while self.child_node.left:
                    self.child_parent_node = self.child_node
                    self.child_node = self.child_node.left

                # 왼쪽 자식 노드 중 가장 작은 값을 가지는 노드에 오른쪽 자식이 있는 경우
                if self.child_node.right != None:
                    self.child_parent_node.left = self.child_node.right
                    
                self.parent_node.right = self.child_node
                self.child_node.left = self.current_node.left
                self.child_node.right = self.current_node.right

        del self.current_node

        return True