이진 탐색 트리(BST) 구현하기(Python3)

승톨·2021년 5월 5일

트리

인트로

최근에 구직을 위해 기술면접을 보면서 화이트보드 코딩 테스트로 이진 탐색 트리 구현을 하게 되었다.

대략적인 트리 구현은 알고 있었지만, 막상 그 자리에서 구현하고, 리뷰하는 것이 만만치 않더라.

(더군다나 화이트보드 코딩 테스트가 처음이라 많이 긴장하기도 했다.😭)

그래서 다시 한번 트리를 정리해서 제대로 소화해보고자 한다. 오늘은 이진 탐색트리(Binary Search Tree, BST) 구현에 대해서 정리해보겠다.

이진 트리 알아보기

본격적인 구현에 앞서 먼저 이진트리부터 알아보자.

이진트리란 '노드가 왼쪽 자식과 오른쪽 자식만을 갖는 트리'이다.

(참고로 "이 글의 독자가 노드, 자식, 부모 등의 트리에 대한 기본적인 개념은 알고 있다."는 전제 아래에 글을 썼다.)

이 트리에선 두 자식 가운데 하나가 존재하지 않아도 상관없다. 또한 두 자식의 대소관계도 상관없다.

참고로 이진트리 중

1. 루트부터 아래쪽 레벨로 노드가 가득 차 있고,
2. 마지막 레벨에서 왼쪽부터 오른쪽으로 노드가 채워져 있는

이진트리를 '완전 이진트리' 라고 한다.

오늘 알아볼 이진 탐색 트리는 무조건 '완전 이진 트리'는 아니다. 2번 조건이 미충족 할 수 있기 때문이다.

이제 이진 탐색 트리를 알아보자.

이진 탐색 트리는 모든 노드가 아래의 조건을 만족해야 한다.

왼쪽 서브트리 노드의 키는 자신의 노드 키보다 작아야 한다.
오른쪽 서브트리 노드의 키는 자신의 노드 키보다 커야 한다.
키가 같은 노드는 복수로 존재하지 않는다.

아래는 이진 탐색 트리를 표현한 그림이다.

노드 5를 보면, 왼쪽 서브트리의 노드들(4,1)은 모두 5보다 작고, 오른쪽 서브트리의 노드(7,6,9)는 모두 5보다 큰 걸 알 수 있다. 이진 탐색 트리의 조건을 충족한다.

이진 탐색 트리 설계

모든 코드는 파이썬3로 작성 되었습니다.

Node 클래스

먼저 이진 탐색 트리의 노드를 나타내는 클래스를 만들어보자.
여기선 4개의 필드로 구성했다.
참고로 value만 넣는게 아니라 key, value pair로 구성했다.
value만 있어도 구현에는 문제 없다. 다만, key-value pair 형태로 가게 되면 value에 다양한 데이터를 넣기가 편하고, key를 정수로 구성하면 조회,삽입,삭제 연산이 상대적으로 간편해져기 때문에 이렇게 구성을 가져갔다.

#트리를 구성하는 노드 클래스
class Node:
    """생성자"""
    def __init__(self, key,value, left,right):
        self.key = key  # 키
        self.value = value # 값
        self.left = left # 왼쪽 자식 참조
        self.right = right # 오른쪽 자식 참조

트리 클래스

트리 클래스는 아래와 같이 구성되어있다.

생성자
노드 검색 메소드
노드 추가 메소드
노드 삭제 메소드
노드 출력 메소드(전위 순회)

class BinarySearchTree():	
  # 생성자
  def __init__(self)->None:

  # 검색하는 메소드
  def search(self, key)->int:
 
  # 노드 추가하는 메소드
  def add(self,key,value)->bool:

  # 노드 삭제하는 메소드
  def remove(self, key)-> bool:
  
	# 노드 출력하는 메소드
  def dump(self) -> None:

이제 하나씩 뜯어보자.

생성자

생성자 로직은 간단하다.
트리의 root를 세팅한다. 초기에는 아무 값도 넣지 않았으니 빈 값으로 둔다.

def __init__(self)->None:
        self.root = None # 루트 설정

검색

기본적인 로직은 다음과 같다.

루트에서 시작하여 찾고자 하는 키의 대소 관계를 판단한 결과에 따라 왼쪽 혹은 오른쪽 서브트리를 따라가면서 검색을 수행한다.

좀 더 구체적인 알고리즘은 다음과 같다.

node라는 루트를 지정한다.
node가 None이면 검색을 실패하고 종료한다.
검색하는 key와 현재 노드 node의 키를 비교한다.
- key == node : 검색을 성공하고 종료한다.
- key < node : 현재 노드를 왼쪽 자식 노드로 옮긴다.
- key > node : 현재 노드를 오른쪽 자식 노드로 옮긴다.
2번으로 돌아가 차례대로 수행을 반복한다.

인자 : key / 리턴값 : 노드 value

def search(self, key)->int:
        node = self.root # 루트에 주목
        while True:
            if node is None: # 더 이상 진행할 수 없으면
                return -1 # 검색 실패
            if key == node.key: # key와 노드 p의 키가 같으면
                return node.value # 검색 성공
            elif key < node.key: # key가 작으면
                node = node.left # 왼쪽 서브 트리에서 검색
            else: # key가 작으면
                node = node.right # 오른쪽 서브 트리에서 검색

노드 추가

노드 추가 시 트리의 형태가 이진 탐색 트리의 조건을 유지해야 한다.

따라서 노드 삽입 시에는 먼저 탐색을 통해 삽입 할 위치를 찾아낸 뒤에 삽입을 수행한다.

알고리즘은 다음과 같다.

root가 있는지 파악한다.
- root가 없으면 생성한다.
- root가 있으면 root부터 시작한다.
(루트가 있다는 전제) 루트 - 현재 노드를 node라고 한다.
삽입하는 key와 현재 노드 node의 키를 비교한다.
- key == node : 추가하려는 키가 이미 있으니 삽입을 실패하고 종료
- key < node :
  - 왼쪽 자식 노드가 없으면, 그 자리에 노드를 만들어서 삽입하고 종료.
  - 왼쪽 자식 노드가 있으면, 현재 노드를 왼쪽 자식 노드로 옮겨서 다시 탐색한다.
- key > node:
  - 오른쪽 자식 노드가 없으면, 그 자리에 노드를 만들어서 삽입하고 종료한다.
  - 오른쪽 자식 노드가 있으면, 현재 노드를 오른쪽 자식 노드로 옮겨서 다시 탐색한다.
3번 과정을 반복한다.

참고로,

루트가 있는 경우 add_node라는 내부 함수를 만들어서 구현했다.
자식 노드로 옮겨가는 과정을 재귀함수 호출로 구현했다.

def add(self,key,value)->bool:
    # 노드 추가하는 내부 함수
    def add_node(node, key,value)->None:
        # 탐색하고 적절한 자리에 삽입
        if key == node.key: #이미 삽입하려는 키가 있으면 false 처리
            return False
        # 삽입하려는 키가 현재 탐색 노드의 키보다 작다면
        elif key < node.key:
            # 그 탐색 노드의 왼쪽 자식이 없다면 바로 그 자리에 삽입
            if node.left is None:
                node.left = Node(key,value,None,None)
            else:
            # 자식이 있으면 재귀함수 호출로 한번 더 들어감
                add_node(node.left,key,value)
        # 삽입하려는 키가 현재 탐색 노드의 키보다 크거나 같다면
        else:
            # 그 탐색 노드의 오른쪽 자식이 없다면 바로 그 자리에 삽입
            if node.right is None:
                node.right = Node(key,value,None,None)
            else:
            # 자식이 있으면 재귀함수 호출로 한번 더 들어감
                add_node(node.right,key,value)
        return True
    # 루트가 있는 경우
    if self.root is None:
        self.root = Node(key,value,None,None)
        return True
    # 루트가 없는 경우
    else: #리턴값은 내부함수 리턴 값
        return add_node(self.root,key,value)

노드 삭제

이번 구현 중에 가장 복잡한 로직이다.

삭제 또한 추가와 마찬가지로 노드의 위치를 찾아낸 뒤에 삭제를 수행해야 한다.

알고리즘은 다음과 같다.

루트 - 현재 노드를 node라고 한다.
삭제하는 key와 현재 노드 node의 키를 비교한다.
- key == node : 삭제 할 키를 찾음. 비교 종료.
- key < node : 현재 노드를 왼쪽 자식 노드로 옮긴다.
- key > node: 현재 노드를 오른쪽 자식 노드로 옮긴다.
삭제할 키를 찾았으면, 케이스를 나눠야 한다.
1. 자식 노드가 없는 노드를 삭제하는 경우
2. 자식 노드가 1개인 노드를 삭제하는 경우
3. 자식 노드가 2개인 노드를 삭제하는 경우
케이스에 따라 삭제를 진행한다.

삭제 시 각각의 케이스를 살펴보자.

자식 노드가 없는 노드를 삭제하는 경우

삭제할 노드가 부모 노드의 왼쪽 자식이면, 부모의 왼쪽 참조를 None으로 한다.
삭제할 노드가 부모 노드의 오른쪽 자식이면, 부모의 오른쪽 참조를 None으로 한다.

자식 노드가 1개인 노드를 삭제하는 경우

삭제 노드(7)가 없어지면 그 자리를 자식 노드(8)가 대체해야 한다.

그러면 자식 노드(8)를 루트로 하는 서브트리의 모든 키가 삭제 노드의 부모노드(6)보다 커지게 되고, 이진 탐색 트리의 조건을 충족한다.

삭제 노드가 부모 노드의 왼쪽 자식인 경우 : 부모의 왼쪽 자식 참조가 삭제할 노드의 자식을 가리키도록 업데이트 한다.
삭제 노드가 부모 노드의 오른쪽 자식인 경우 : 부모의 오른쪽 자식 참조가 삭제할 노드의 자식을 가리키도록 업데이트 한다.

여기까지의 코드를 한번 살펴보자.

위의 로직을 구현하기 위해 parent를 세팅했고, 현재 노드는 node로 세팅했다.

그리고 삭제 시 자식 노드의 방향을 한번에 판단하기 위해 is_left_child 라는 플래그를 세팅했다.

def remove(self, key)-> bool:
    node = self.root # 현재 노드로 지정
    parent = None  # 현재 노드의 부모 노드
    is_left_child = True # node는 parent의 왼쪽 자식 노드인지 오른쪽 자식 노드인지 확인

    # 삭제할 노드 탐색
    while True:
      if node is None:
          return False
      if key == node.key:
          break
      else:
          parent = node
          if key < node.key:
              node = node.left
              is_left_child = True # 왼쪽 자식 노드로 내려가니까 플래그를 True로 설정
          else:
              node = node.right
              is_left_child = False # 오른쪽 자식 노드로 내려가니까 플래그를 True로 설정
    
    # 키를 찾은 뒤에 자식이 없는 노드이면 or 자식이 1개 있는 노드이면
    if node.left is None: # 왼쪽 자식이 없으면
        if node is self.root: #만약 삭제 노드가 root이면, 바로 오른쪽 자식으로 대체한다.
            self.root = node.right
        # 아래의 parent는 탐색 시 찾은 노드의 바로 위 부모가 됨.(탐색 로직에서 그렇게 적용)
        # parent - node - node의 자식의 구도가 있으면 node라는 중간이 빠지기 때문에 parent의 자식과 node의 자식을 연결
	# (node의 자식이 없으면 자연스레 None이 들어감)
	elif is_left_child: #왼쪽 자식 노드가 있는 것이니까
            parent.left = node.right # 부모의 왼쪽 참조가 오른쪽 자식을 가리킴
        else: #오른쪽 자식 노드가 있는 것이니까
            parent.right = node.right # 부모의 오른쪽 참조가 오른쪽 자식을 가리킴
    
    elif node.right is None: # 오른쪽 자식이 없으면
        if node is self.root: 
            self.root = node.left #만약 삭제 노드가 root이면, 바로 왼쪽 자식으로 대체한다.
        elif is_left_child:
            parent.left = node.left # 부모의 왼쪽 참조가 왼쪽 자식을 가리킴
        else:
            parent.right = node.left # 부모의 오른쪽 참조가 왼쪽 자식을 가리킴

이제 자식 노드가 2개인 경우를 살펴보자.

자식 노드가 2개인 노드를 삭제하는 경우

이 로직에서는 무조건 왼쪽 서브트리에서 대체할 노드를 찾으려고 한다. 왼쪽 서브트리에서 가장 큰 노드로 대체할 것이다. 참고로, 어느쪽에서 대체노드를 찾든지 상관은 없다.

삭제할 노드의 왼쪽 서브트리에서 키 값이 가장 큰 노드를 검색한다.
검색 한 노드를 삭제 할 노드 위치로 옮긴다. (데이터 복사)
옮긴 노드를 삭제한다. 삭제 시 자식 노드의 개수에 따라 데이터가 조정된다.
- 옮긴 노드에 자식이 없으면 : 옮긴 노드의 부모는 자식으로 None을 가진다.
- 옮긴 노드에 자식이 1개 있으면 : 옮긴 노드의 부모는 그 자식으로 옮긴 노드의 자식을 가진다.

삭제 로직은 특히 글만 보면 이해가 잘 안 될 수 있다. 그래서 그림을 준비해보았다.

아래와 같은 트리가 있다고 가정해보자.

(10→11→12→6→2→8→7→9→1→4→3→5로 삽입)

이 트리에서 6을 없애보자.

로직은 아래와 같다.

삭제하려는 노드 6을 찾는다.
6의 자식 노드는 2개이다.
parent를 설정하고, 왼쪽 서브 트리 중 가장 큰 노드를 찾기 위해 node_max_left를 설정한다.
가장 큰 노드를 찾기 위해 탐색한다.
그 노드(node_max_left)를 찾았으면 그 노드의 key, value를 삭제하려는 노드에 넣는다.(대체한다.)
node_max_left를 제거한다.(해당 노드의 부모 노드와 해당 노드의 자식을 연결한다. 자식노드가 없으면 None으로 설정)

아래 그림은 삭제 하는 과정을 묘사한 그림이다.

회색 노드가 삭제 노드이고, 빨간색 테두리를 가진 노드가 node_max_left가 된다.
node_max_left가 회색 노드를 대체한다.
node_max_left가 제거될 때 node_max_left.left(None)를 부모노드(4)와 연결한다.

삭제가 완료되면 아래와 같은 트리가 나온다.

이제 5를 지워보자.

이번에도 자식노드가 2개이므로, 위의 삭제로직과 동일하다.

node_max_left는 4가 된다. 다만 이번에는 node_max_left에 왼쪽 자식(3)이 있다. 따라서 부모노드(2)를 왼쪽 자식과 연결시킨다.

이번에는 노드 8을 지워본다. 8의 서브트리 중 가장 큰 노드 7로 대체되어서 아래와 같은 그림이 된다.

이제 7을 지워보자.

노드 7의 자식 노드는 1개이므로, 위의 '자식 노드가 1개인 노드를 삭제하는 경우'가 적용된다.

또한 자식 노드는 오른쪽 자식 노드이다. 그 부분만 다시 살펴보자.

  # 키를 찾은 뒤에 자식이 없는 노드이면 or 자식이 1개 있는 노드이면
    if node.left is None: # 왼쪽 자식이 없으면
        if node is self.root: #만약 삭제 노드가 root이면, 바로 오른쪽 자식으로 대체한다.
            self.root = node.right
        # 아래의 parent는 탐색 시 찾은 노드의 바로 위 부모가 됨.(탐색 로직에서 그렇게 적용)
        # parent - node - node의 자식의 구도가 있으면 node라는 중간이 빠지기 때문에 parent의 자식과 node의 자식을 연결
				# (node의 자식이 없으면 자연스레 None이 들어감)
				elif is_left_child: #왼쪽 자식 노드가 있는 것이니까
            parent.left = node.right # 부모의 왼쪽 참조가 오른쪽 자식을 가리킴
        else: #오른쪽 자식 노드가 있는 것이니까
            parent.right = node.right # 부모의 오른쪽 참조가 오른쪽 자식을 가리킴

노드 7이 지워지게 되면, 이진 탐색 트리 조건을 유지하기 위해 부모 노드와 자식 노드를 연결해야 한다.
노드 7은 부모 노드(5)의 오른쪽 자식이기 때문에 is_left_child 는 False인 상태이다.
따라서 부모 노드의 오른쪽 자식은 노드 7의 오른쪽 자식(9)이 된다.

최종 트리는 아래와 같은 그림이 된다.

이렇게 트리의 삭제를 그림으로 살펴보았다.

자식노드가 2개일 때 삭제하는 경우는 코드를 보면서 한번 더 이해해보자.

# 자식이 2개 있는 노드이면
else: # 무조건 왼쪽 서브트리에서 대체할 노드를 찾는다. 왼쪽 서브트리에서 가장 큰 노드로 대체한다.
    parent = node
# parent를 지정한 이유 : 지우는 노드로 지정되는데, 하위 자식 노드가 많으면 node 삭제 시 조정이 일어남. 
# node말고 그 하위 노드들을 관리할 주체가 필요함. 
# 그때 가장 하단의 자식 노드의 연결을 끊기위해 parent를 일단 삭제할 node로 지정.
    node_max_left = node.left # 왼쪽 서브트리에서 가장 큰 노드를 찾기 위해 초기값 설정
    is_left_child = True # 왼쪽 서브트리에서 가장 큰 노드가 부모 노드와 어떤 관계(왼쪽,오른쪽)인지 파악하기 위해 세팅
    
# 가장 큰 노드를 찾기 위해 탐색
    while node_max_left.right is not None:
        parent = node_max_left
        node_max_left = node_max_left.right
        is_left_child = False

    # 가장 큰 노드를 찾았으니 삭제하려는 노드를 대체
    node.key = node_max_left.key
    node.value = node_max_left.value

    # is_left_child가 트루가 되려면, 삭제 노드의 왼쪽 서브트리중 오른쪽 자식이 없어야 함. 
    if is_left_child:
# 무조건 node_max_left.left로 지정하는 이유 : 
# 1. 가장 큰 노드가 왼쪽 자식을 갖고 있을 수 있음.
# 2. 이미 오른쪽 자식 노드는 위에서 탐색을 완료했기 때문에 여기서 추가적인 오른쪽 자식 노드를 지정할 수 없다.
# 3. 그러므로 삭제 노드의 left를 node_max_left의 left로 연결(자식이 없으면 None을 가지게 됨.) 
# node_max_left의 왼쪽 자식 노드만 있거나 자식이 없는 경우만 가능. 자식이 없으면 None으로 적용됨.
        parent.left = node_max_left.left
    else:
        parent.right = node_max_left.left
return True # 정상적으로 조정되었으니 true

노드 출력

모든 노드의 key, value를 출력한다.
전위 순회 알고리즘을 활용한 깊이 우선 탐색을 적용하였다.

참고 : 전위순회, 중위순회, 후위 순회에 대해 궁금하다면?
http://ejklike.github.io/2018/01/09/traversing-a-binary-tree-2.html

구현을 용이하게 하기 위해 내부 함수 print_subtree()를 만들었다.

# 노드 출력하는 메소드
def dump(self):
    def print_subtree(node):
        # 전위 순회로 출력
        if node is not None:
            print(f'{node.key} {node.value}')
            print_subtree(node.left)
            print_subtree(node.right)
    root = self.root
    print_subtree(root)