[좌충우돌, 파이썬으로 자료구조 구현하기] 이진 탐색 트리 구현하기

숫자 맞추기(Up Down) 게임은 이진 검색을 이용한 것이다.

1 ~ 100 사이에 임의의 숫자를 맞추는 놀이를 할 때, 처음 추측하는 값은 50으로 한다. 그러면 이 값보다 크거나 작은지에 따라 Up, Down이라고 얘기하면서 다시 그 사이의 값을 추측하는 방식으로 답을 찾는다.

이러한 방식이 이진 검색이다. Up과 Down을 왼쪽과 오른쪽으로 바꿔서 생각해보자.

이진 트리는 앞서 공부했으므로 이진 트리와 이진 탐색 트리의 차이점을 알아보자.

• 모든 왼쪽 서브 트리의 노드는 루트 노드보다 작다.
• 모든 오른쪽 서브 트리의 노드는 루트 노드보다 크다.
• 왼쪽과 오른쪽 서브 트리도 모두 이진 탐색 트리이다.
• 중복 노드는 없다.

(a)는 이진 탐색 트리이지만, (b)는 이진 탐색 트리가 아니다. 값이 1인 노드는 4의 오른쪽에 있으므로 4보다 커야 하는데, 이 조건을 만족하지 않으므로 이진 탐색 트리가 아니다.

Binary Search Tree 클래스 구현하기

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

class BS_Tree:
    def __init__(self):
        self.root = None

1. contains 메서드 만들기

특정 값을 찾는 메서드를 만든다. __contains__ 메서드로 만들면, 다른 컨테이너형 객체처럼 in 또는 not in을 사용할 수 있다.

class Node:
    ...

class BS_Tree:
    ...

    def __contains__(self,data):
        node = self.root
        while node:
            if node.data == data:
                return True
            elif node.data > data:
                node = node.left
            else:
                node = node.right
        return False

• 루트 노드부터 data와 비교하여 값이 같으면 True를 반환한다.
• 매개변수로 주어진 data가 노드 data보다 작으면 왼쪽 노드로 이동하여 비교한다.

2. insert 메서드 구현하기

이진 탐색 트리를 만족하도록 key,data를 삽입하는 메서드를 만든다.

2-1. 반복문으로 구현하기

def insert(self,data):
    if self.root is None:
        self.root = Node(data)
        return

    node = self.root
    while True:
        if node.data > data:
            if node.left is None:
                node.left = Node(data)
                return
            else:
                node = node.left
        else:
            if node.right is None:
            node.right = Node(data)
            return
        else:
            node = node.right

2-2. 재귀 함수로 구현하기

• 앞선 반복문 코드가 if...else 구문도 많고 반복되는 부분도 있어 재귀로 구현해본다.
• 루트 노드의 서브 트리 노드들도 이진 탐색 트리이기 때문에 재귀적으로 접근할 수 있다.

def insert_recursive(self,data):
	def _insert(node, data):
            if node is None:
                return Node(data)

            if node.data > data:
                node.left = _insert(node.left, data)
            else:
                node.right = _insert(node.right, data)
            return Node
        
	self.root = _insert(self.root, data)

3. delete 메서드 구현하기

삭제할 노드를 찾고 삭제할 노드의 연결 상태에 따라 이진 탐색 트리를 재구성해야한다.

 

삭제할 노드의 경우의 수는 다음과 같다.

 

1. 리프 노드일 경우: 그냥 삭제

2. 자식 노드가 하나일 경우: 자식 노드를 현재 노드 위치로 끌어올린다.

3. 자식 노드가 2개일 경우

• 오른쪽 자식 노드 중에서 가장 작은 노드를 끌어올린다. 혹은 왼쪽 자식 노드 중에서 가장 큰 노드를 끌어올린다.
• 오른쪽 자식 노드가 서브 트리면, 서브 트리에서 가장 왼쪽 노드가 가장 작은 노드이다.

(a) 삭제할 노드가 리프 노드일 경우

• 값이 2,5,8 인 노드는 리프 노드이다. 이 노드를 삭제하려면, 해당 노드는 None으로 만든 후 부모 노드에 연결한다.

(b) 삭제할 노드의 자식이 하나일 경우

• 값이 7, 9 노드의 자식 노드는 왼쪽 또는 오른쪽 하나만 있다. 예를 들어 값이 7인 노드를 삭제하려면 삭제할 노드와 삭제할 노드의 부모 노드를 찾는다.
• 값이 7인 노드의 자식 노드(8)를 부모 노드(9)의 왼쪽에 연결한다.

9를 삭제할 때도 위와 같이 처리한다.

(c) 삭제할 노드의 자식이 2개인 경우

• 삭제할 노드와 삭제할 노드의 부모 노드를 찾는다.
• 값이 4인 노드의 오른쪽 자식 중에서 값이 가장 작은 노드(5)를 찾는다.
  • 즉, 값이 가장 작은 노드는 리프 노드이거나 오른쪽 자식만 가진다.
• 삭제할 노드의 값을 위에서 찾은 가장 작은 노드의 값(5)로 대체한다.
• 값이 5인 노드를 삭제한다.
  • 이는 (a)인 경우와 동일하다.
  • 만약 5의 오른쪽 자식이 있다면 (b)와 동일하다.

# 삭제할 노드가 리프 노드이거나, 자식이 하나일 때 처리한다.
    # 왼쪽 자식이 있으면 반환하고, 그렇지 않으면 오른쪽 자식을 반환한다.
    # 리프 노드이면 왼쪽과 오른쪽이 모두 None이므로 None을 반환한다.
    def _del_node1(self,del_node):
        if del_node.left:
            return del_node.left
        return del_node.right
    
    # 자식이 둘인 노드를 삭제한다.
    def _del_node2(self,del_node):
        # 삭제할 노드의 부모 노드와 오른쪽에서 가장 왼쪽 노드를 찾는다.
        parent = del_node
        leftmost = del_node.right
        while leftmost.left:
            parent = leftmost
            leftmost = leftmost.left
        
        # 삭제할 노드의 값을 가장 왼쪽 노드의 값으로 대체한다.
        del_node.data = leftmost.data
        
        # 가장 왼쪽 노드를 삭제한다.
        # 가장 왼쪽 노드는 리프 노드이거나 오른쪽에 자식이 있다.
        parent.left = self._del_node1(leftmost)
        
    def delete(self,data):
        # 루트가 None일 때와 삭제 대상일 때를 처리한다.
        if self.root is None:
            return
        if self.root.data == data:
            if self.root.left and self.root.right:
                self._del_node2(self.root)
            else:
                self.root = self._del_node1(self.root)
            return
        
        #삭제할 노드와 그 부모 노드를 찾는다.
        del_node = self.root
        while del_node and del_node.data != data:
            parent = del_node
            if del_node.data > data:
                del_node = del_node.left
            else:
                del_node = del_node.right

        if del_node is None: #삭제할 노드가 없으면 종료
            return

        #삭제할 노드의 자식이 둘일 때
        if del_node.left and del_node.right:
            self._del_node2(del_node)
            return
        #삭제할 노드가 리프 노드이거나 자식이 하나일 때
        if parent.left == del_node:
            parent.left = self._del_node1(del_node)
        else:
            parent.right = self._del_node1(del_node)