자료구조 트리 구현(파이썬)

트리 구조는 구현하기 힘들지만 이진탐색 알고리즘이나 힙을 통한 우선순위 큐 등 다양한 분야로 응용 될 수 있다.

때문에 한번쯤은 구현을 해보는 게 좋다고 생각했다.

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

class NodeMgmt:
    def __init__(self, head):
        self.head = head

    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.current_node = self.head
        while self.current_node:
            if self.current_node.value == value:
                return True
            elif value < self.current_node.value:
                self.current_node = self.current_node.left
            else:
                self.current_node = self.current_node.right
        return False

    def delete(self, value):
        searched = False
        self.current_node = self.head
        self.parent = self.head
        while self.current_node:
            if self.current_node.value == value:
                searched = True
                break
            elif value < self.current_node.value:
                self.current_node = self.current_node.left
            else:
                self.current_node = self.current_node.right

        # value를 찾았을 경우
        if not searched:
            return False
        else:
        #case1  child 노드가 없는 노드 삭제 경우
            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
                del self.current_node
        #case2 child 노드가 하나 있는 삭제 경우

            if self.current_node.left != None and self.current_node.right == None:
                # case2-1 child 노드가 하나 있는데 왼쪽일 경우
                if value < self.parent.value:
                    # 지우려는 노드가 부모 노드의 왼쪽에 있는 경우
                    self.parent.left = self.current_node.left
                    # 현재 노드의 왼쪽 노드를 부모의 왼쪽에 붙여줌
                else:
                    # 지우려는 노드가 부모 노드의 오른쪽에 있는 경우
                    self.parent.right = self.current_node.left
                    # 현재 노드의 왼쪽 노드를 부모의 오른쪽에 붙여줌

            elif self.current_node.left == None and self.current_node.right != None:
                #case2-2 child 노드가 하나 있는데 오른쪽일 경우
                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:
                # case3 child 노드가 두 개 모두가 있을 경우
                if value < self.parent.left:
                    # case3-1 삭제할 노드(current_node)가 parent의 왼쪽에 있는 경우
                    self.change_node = self.current_node.right
                    self.change_node_parent = self.parent
                    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
                        # 오른쪽 노드가 없으면 그냥 비워 버림
                    self.parent.right = self.change_node
                    # 삭제 노드의 부모의 오른쪽에 바꿀 노드를 갖다붙임
                    self.change_node.left = self.current_node.left
                    # 바꿀 노드의 왼쪽에 삭제 노드의 왼쪽을 갖다 붙이고
                    self.change_node.right = self.current_node.right
                    # 바꿀 노드의 오른쪽에 삭제노드의 오른쪽을 갖다 붙임
        return True

삭제하는 부분이 어렵다. 여러가지 케이스가 있기 때문에...

하지만 발생하는 케이스들을 꼼꼼하게 정리해서 분할하면 if 문을 통해 쪼개서 해결할 수 있다.