💻 백준 1991번: 트리 순회

문제 소개 🧐

이진 트리를 입력받아 전위 순회(preorder traversal), 중위 순회(inorder traversal), 후위 순회(postorder traversal)한 결과를 출력하는 프로그램을 작성하는 문제였다.

입력 📝

첫째 줄에 이진 트리의 노드 개수 N(1 ≤ N ≤ 26)이 주어진다.
둘째 줄부터 N개의 줄에 걸쳐 각 노드와 그의 왼쪽 자식 노드, 오른쪽 자식 노드가 주어진다.
노드의 이름은 A부터 차례대로 알파벳 대문자로 매겨지며, 항상 A가 루트 노드가 된다. 자식 노드가 없는 경우에는 .으로 표현한다.

출력 📤

총 3줄에 걸쳐 전위 순회, 중위 순회, 후위 순회한 결과를 각각 출력한다. 각 줄에 N개의 알파벳을 공백 없이 출력하면 된다.

제한 ❌

  • 시간: 2초
  • 메모리: 128 MB

나의 풀이: 반복문과 큐를 이용한 시도 😥

Idea

처음엔 Node 클래스를 직접 정의하고, 각 노드를 리스트에 저장한 뒤 반복문과 큐(deque)를 사용해 순회를 구현하려 했다. 전위 순회는 스택처럼, 나머지는 큐처럼 동작하게 하려다 보니 로직이 매우 복잡해졌다.

Code

import sys
from collections import deque

class Node:
    val: int
    left: "Node"
    right: "Node"
    parent: "Node"
    visited: bool

    def __init__(self, val: str):
        self.val = val
        self.left = None
        self.right = None
        self.parent = None
        self.visited = False

    def set_child(self, left, right):
        self.left = left
        self.right = right
    
    def set_parent(self, parent):
        self.parent = parent

    def __str__(self):
        return self.val

    def get_right(self):
        return self.right
    def get_left(self):
        return self.left
    def get_parent(self):
        return self.parent

n = int(sys.stdin.readline().rstrip())
node_list = list()
in_list = [list(map(lambda x: ord(x)-65, sys.stdin.readline().split())) for _ in range(n)]

for i in range(65, 91):
    node = Node(chr(i))
    node_list.append(node)

def reset_node_tree(in_list):
    global node_list
    for li in in_list:  # '.' 은 46번
        node_idx, l_idx, r_idx = li
        node: Node = node_list[node_idx]
        left = node_list[l_idx] if 0 <= l_idx else None
        right = node_list[r_idx] if 0 <= r_idx else None
        node.set_child(left, right)
        left.set_parent(node) if left != None else None
        right.set_parent(node) if right != None else None
    for node in node_list:
        node.visited = False

# 전위 순회 v l r
def prev_search():
    global node_list
    q = deque()
    q.append(node_list[0])
    answer = ""
    while q:
        node = q.pop()
        answer += node.__str__()
        if node.get_right() is not None:
            q.append(node.get_right())
        if node.left is not None:
            q.append(node.get_left())
    print(answer)

def mid_search(): # 중위 순회, l v r
    global node_list
    q = deque()
    q.append(node_list[0])
    answer = ""
    while q:
        node = q.popleft()
        if node.left != None:
            q.appendleft(node.get_left())
        elif not node.visited:
            answer += node.__str__()
            parent = node.get_parent()
            node.visited = True
            if parent != None:
                parent.set_child(None, None)
                q.appendleft(parent)

        if node.right != None:
            q.append(node.get_right())
            
    print(answer)

def end_search(): # 후위 순회, l r v
    global node_list
    q = deque()
    q.append(node_list[0])
    answer = ""
    while q:
        node = q.popleft()
        # 자식 노드 담기
        if node.left != None:
            q.append(node.get_left())
        if node.right != None:
            q.append(node.get_right())
        
        if node.left is None and node.right is None and not node.visited:  # 말단 노드면
            parent = node.get_parent()
            if parent is None:
                continue
            if node == parent.left: # parent와의 연결을 끊음
                parent.set_child(None, parent.right)
            else:
                parent.set_child(parent.left, None)
            q.appendleft(parent)
            node.visited = True
            answer += node.__str__()
            
    print(answer + 'A')

reset_node_tree(in_list)
prev_search()
mid_search()
reset_node_tree(in_list)
end_search()

Result

실패! 😱

중위, 후위 순회를 반복문으로 구현하는 로직이 너무 복잡했다. 특히 방문한 노드를 처리하고 부모 노드로 돌아가는 과정에서 예외 케이스가 많아 결국 해결하지 못했다.

두 번째 시도: 재귀를 활용한 풀이 🎉

Idea ✨

복잡한 반복문 대신 재귀(Recursion) 를 사용하면 훨씬 간결하게 트리 순회를 구현할 수 있다는 것을 깨달았다.

  • 전위 순회: (루트) -> (왼쪽) -> (오른쪽) 순서로 재귀 호출
  • 중위 순회: (왼쪽) -> (루트) -> (오른쪽) 순서로 재귀 호출
  • 후위 순회: (왼쪽) -> (오른쪽) -> (루트) 순서로 재귀 호출

또한, 노드를 리스트가 아닌 딕셔너리로 관리하면 {'A': Node('A'), ...} 와 같이 각 노드에 이름으로 바로 접근할 수 있어 편리하겠다고 생각했다.

Code

import sys

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

def preorder(node):
    if node is None:
        return ""
    return node.val + preorder(node.left) + preorder(node.right)

def inorder(node):
    if node is None:
        return ""
    return inorder(node.left) + node.val + inorder(node.right)

def postorder(node):
    if node is None:
        return ""
    return postorder(node.left) + postorder(node.right) + node.val

n = int(sys.stdin.readline().rstrip())
nodes = {chr(65 + i): Node(chr(65 + i)) for i in range(26)}  # A~Z 미리 생성

for _ in range(n):
    p, l, r = sys.stdin.readline().split()
    if l != '.':
        nodes[p].left = nodes[l]
    if r != '.':
        nodes[p].right = nodes[r]

root = nodes['A']  # 항상 A가 루트라고 가정

print(preorder(root))
print(inorder(root))
print(postorder(root))

Result

성공! ✅

재귀를 사용하니 코드가 놀랍도록 간결해지고, 각 순회의 논리가 명확하게 드러났다. 딕셔너리로 노드를 관리하니 특정 노드를 찾아 연결하는 과정도 훨씬 쉬워졌다.

배운 점 및 느낀 점 ✍️

  • 재귀의 힘: 트리, 그래프와 같이 재귀적인 자료구조는 반복문보다 재귀 함수를 사용하는 것이 훨씬 직관적이고 효율적일 수 있다는 것을 깨달았다.
  • 자료구조 선택의 중요성: 문제의 특성에 맞는 적절한 자료구조(이번엔 딕셔너리)를 선택하는 것이 코드의 간결성과 효율성을 크게 좌우한다는 것을 배웠다.
  • 실패로부터 배우기: 첫 번째 시도의 실패 원인을 분석하고 다른 접근법을 고민하는 과정에서 더 나은 해결책을 찾을 수 있었다. 복잡한 코드는 버리고 더 간단한 해결책을 찾는 용기도 필요하다.
  • 기본기: 트리 순회와 같은 기본적인 개념에 대한 확실한 이해가 뒷받침되어야 응용 문제를 해결할 수 있다.