[WEEK03] 백준 1991 트리 순회

문제

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

예를 들어 위와 같은 이진 트리가 입력되면,

전위 순회한 결과 : ABDCEFG // (루트) (왼쪽 자식) (오른쪽 자식)
중위 순회한 결과 : DBAECFG // (왼쪽 자식) (루트) (오른쪽 자식)
후위 순회한 결과 : DBEGFCA // (왼쪽 자식) (오른쪽 자식) (루트)
가 된다.

입력

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

출력

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

풀이

2가지 방법으로 풀었다.

첫번째 방법은 딕셔너리를 사용 입력받은 값을 키값으로 해당 노드를 출력해주었다.

import sys
input = sys.stdin.readline

def PreOrder(root):
    if root != '.':
        print(root, end='')
        PreOrder(a[root][0])
        PreOrder(a[root][1])

def InOrder(root):
    if root != '.':
        InOrder(a[root][0])
        print(root, end='')
        InOrder(a[root][1])

def PostOrder(root):
    if root != '.':
        PostOrder(a[root][0])
        PostOrder(a[root][1])
        print(root, end='')

n = int(input())
a = {}

for _ in range(n):
    root, left, right = sys.stdin.readline().split()
    a[root] = [left, right]

PreOrder('A')
print()
InOrder('A')
print()
PostOrder('A')

두번째 방법은 직접 Node, BST 클래스를 만들어 주었다.
해당 문제를 풀 수 있을정도의 기능만 구현해 놓았다.

import sys
input = sys.stdin.readline

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

    def link(self, left, right):
        self.left = left
        self.right = right

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

    def search(self, key):
        node = self.root
        
        # 루트가 존재하지 않으면 None
        if not node:
            return None
        # 현재 노드의 키값이 찾고자 하는 키값과 일치하면 현재 노드 반환
        if key == node.key:
            return node
        
        # 루트를 제외한 모든 곳을 탐색한다
        stack = []
        if node.left:
            stack.append(node.left)
        if node.right:
            stack.append(node.right)

        # 모든 곳을 탐색하다 찾는 키값이 나오면 해당 노드를 반환
        while stack:
            node = stack.pop()

            if node.key == key:
                return node
            
            if node.left:
                stack.append(node.left)
            if node.right:
                stack.append(node.right)

        # 여기까지 오면 찾는 노드가 어디에도 없다는거임
        return None
            
    def add(self, p_key, cl_key, cr_key):
        if not self.root:
            self.root = Node(p_key)

        # cl_key와 cr_key는 p_key를 가지는 노드의 자식임
        # p_key를 가지는 노드를 반환 받음
        # cl_key를 가지는 노드를 반환 받은 노드의 left에 연결
        # cr_key를 가지는 노드를 반환 받은 노드의 right에 연결
        if cl_key != '.':
            self.search(p_key).left = Node(cl_key)
        if cr_key != '.':
            self.search(p_key).right = Node(cr_key)

    def pre_order(self, cur: Node):
        print(cur.key, end = '')
        if cur.left:
            self.pre_order(cur.left)
        if cur.right:
            self.pre_order(cur.right)

    def in_order(self, cur: Node):
        if cur.left:
            self.in_order(cur.left)
        print(cur.key, end = '')
        if cur.right:
            self.in_order(cur.right)

    def post_order(self, cur: Node):
        if cur.left:
            self.post_order(cur.left)
        if cur.right:
            self.post_order(cur.right)
        print(cur.key, end = '')


n = int(input())
bst = BST()

for i in range(n):
    p_key, cl_key, cr_key = input().split()
    bst.add(p_key, cl_key, cr_key)

bst.pre_order(bst.root)
print()
bst.in_order(bst.root)
print()
bst.post_order(bst.root)

첫 줄 입력에서 'A'를 키로 가지는 노드가 만들어짐과 동시에 루트가 된다.
이 후에 left, right 노드는 p_key 값을 가진 노드를 반환 받아 해당 노드에 left, right에 각각 연결해주었다.