(Python) 백준 5639

백준 5639

시간 제한	메모리 제한	제출	정답	맞힌 사람	정답 비율
1 초	256 MB	40413	15860	11240	38.247%

문제

이진 검색 트리는 다음과 같은 세 가지 조건을 만족하는 이진 트리이다.

• 노드의 왼쪽 서브트리에 있는 모든 노드의 키는 노드의 키보다 작다.
• 노드의 오른쪽 서브트리에 있는 모든 노드의 키는 노드의 키보다 크다.
• 왼쪽, 오른쪽 서브트리도 이진 검색 트리이다.

전위 순회 (루트-왼쪽-오른쪽)은 루트를 방문하고, 왼쪽 서브트리, 오른쪽 서브 트리를 순서대로 방문하면서 노드의 키를 출력한다. 후위 순회 (왼쪽-오른쪽-루트)는 왼쪽 서브트리, 오른쪽 서브트리, 루트 노드 순서대로 키를 출력한다. 예를 들어, 위의 이진 검색 트리의 전위 순회 결과는 50 30 24 5 28 45 98 52 60 이고, 후위 순회 결과는 5 28 24 45 30 60 52 98 50 이다.

이진 검색 트리를 전위 순회한 결과가 주어졌을 때, 이 트리를 후위 순회한 결과를 구하는 프로그램을 작성하시오.

입력

트리를 전위 순회한 결과가 주어진다. 노드에 들어있는 키의 값은 106보다 작은 양의 정수이다. 모든 값은 한 줄에 하나씩 주어지며, 노드의 수는 10,000개 이하이다. 같은 키를 가지는 노드는 없다.

출력

입력으로 주어진 이진 검색 트리를 후위 순회한 결과를 한 줄에 하나씩 출력한다.

---

답안

import sys

def append(value, bst):
    if len(bst) == 0:
        bst.extend([value, [], []])
        return
    if value < bst[0]:
        append(value, bst[1])
    elif value > bst[0]:
        append(value, bst[2])

def print_post_order(bst):
    if len(bst) == 0:
        return
    print_post_order(bst[1])
    print_post_order(bst[2])
    print(bst[0])

sys.setrecursionlimit(20_000)
graph = []
for num in sys.stdin:
    append(int(num), graph)
print_post_order(graph)

주의: 이 답안은 Python3 구현체로 채점하면 시간 초과 판정을 받습니다. PyPy3로 채점해야 정답 판정을 받을 수 있습니다.

풀이

전위 순회를 통해 BST를 만들어 그 BST를 후위 순회를 하는 문제이다.

이 문제에서 가장 중요한 것은 BST를 담는 자료구조라고 생각하는데, 내가 고려했던 선택지는 다음과 같다.

• Node와 Tree 객체를 만들어 구현 → 객체를 다루는 데에 시간이 많이 들 것 같아 하지 않기로 함
• 1차원 배열을 만들어 트리를 저장 (깊이 0(루트): 인덱스 0, 깊이 1: 인덱스 1~2, 깊이 3: 인덱스 3~6, …) → 이 문제에서 주어지는 최대 노드의 개수는 1만 개이다. 이때 필요한 배열은 $2^0+2^1+2^2+...+2^{10000}$개이므로, 메모리를 과도하게 많이 쓰게 되어 하지 않기로 함

마지막으로 생각했던 방법은 파이썬의 리스트를 다음과 같이 이용하는 것이었다.

트리 = [루트, [왼쪽 서브트리], [오른쪽 서브트리]]

이 구조를 이용하면 시간과 메모리를 적절하게 사용할 수 있을 것 같아 이렇게 하기로 정했다.

def append(value, bst):
    if len(bst) == 0:
        bst.extend([value, [], []])
        return
    if value < bst[0]:
        append(value, bst[1])
    elif value > bst[0]:
        append(value, bst[2])

...

graph = []
for num in sys.stdin:
    append(int(num), graph)

먼저 그래프를 초기화해주는 부분이다.

BST의 정의에 맞게, 입력을 하나씩 받으며 초기화해준다.

def print_post_order(bst):
    if len(bst) == 0:
        return
    print_post_order(bst[1])
    print_post_order(bst[2])
    print(bst[0])

후위순회는 재귀를 통해 구현했다.

leaf node의 left와 right에는 []와 같이 빈 배열이 담기게 되므로, 이 경우 재귀를 종료한다.