백준 5639번 : 이진 검색 트리 in Python

5639번: 이진 검색 트리

 

www.acmicpc.net

이 문제는 이진 검색 트리를 전위 순회한 결과가 주어졌을 때 해당 트리의 후위 순회한 결과를 출력하는 문제이다.

이전의 트리의 순회 문제의 경우 중위 순회와 후위 순회가 주어져야 트리의 구조를 정확히 파악할 수 있었다. 반면에 이진 검색 트리의 경우 특별한 조건 하에 트리가 만들어지기 때문에 전위 순회한 결과만 주어졌을 때 트리의 구조를 온전히 파악할 수 있다.

이진 검색 트리란 기존의 이진 트리에서 검색을 편하게 하기 위해 특정 조건을 추가한 트리로, 각 노드에 대해 노드의 왼쪽 서브트리에 있는 모든 노드의 키는 노드의 키보다 작고, 오른쪽 서브트리에 있는 모든 노드의 키는 노드의 키보다 커야 한다. 

그렇다면 이러한 조건을 이용해 전위 순회를 이용해 어떻게 후위 순회 결과를 출력할 수 있을까? 우선 루트는 전위 순회에서 가장 앞에 있는 노드이다. 왜냐하면 전위 순회는 루트 -> 왼쪽 서브트리 -> 오른쪽 서브트리 순서대로 순회하기 때문이다. 문제는 왼쪽 서브트리와 오른쪽 서브트리를 파악하는 것인데, 이를 파악할 때 이진 검색 트리의 특징을 이용한다. 왼쪽 서브트리의 모든 노드의 키는 루트 노드의 키보다 작고, 오른쪽 서브트리의 모든 노드의 키는 루트 노드의 키보다 크므로, 루트 노드 뒤의 모든 노드들을 탐색하면서 루트 노드의 키보다 큰 키를 가지는 첫번째 노드를 찾는다. 이 첫번째 노드의 위치가 바로 왼쪽 서브트리의 노드들과 오른쪽 서브트리의 노드들의 경계가 된다.

이를 파악할 수 있다면 이제 후위 순회를 출력하는 함수를 만들 수 있다. 우선 함수는 트리를 나타내는 배열의 시작 인덱스와 끝 인덱스를 매개변수로 갖는다. 이후 함수 내에서 위와 같은 방식으로 왼쪽 서브트리와 오른쪽 서브트리의 경계를 찾는다. 그리고 왼쪽 서브트리와 오른쪽 서브트리에 대해 함수를 재귀적으로 호출한 뒤 루트 노드의 키를 출력한다.

이제 이 함수에 시작 인덱스로 0, 끝 인덱스로 트리의 크기-1을 넣으면 후위 순회 결과를 출력할 수 있다.

이를 코드로 나타내면 다음과 같다.

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**5)

# 후위 순회 출력 함수
def post_order(start, end):
    if start > end: return
    # 왼쪽 서브트리와 오른쪽 서브트리의 경계
    division = end+1
    for i in range(start+1, end+1):
        if pre[i] > pre[start]:
            division = i
            break
    # 왼쪽 서브트리에 대해 함수를 재귀적으로 호출
    post_order(start+1, division-1)
    # 오른쪽 서브트리에 대해 함수를 재귀적으로 호출
    post_order(division, end)
    # 루트 노드의 키 출력
    print(pre[start])

pre = []
count = 0
# 입력값을 받기 위해 예외문을 사용한다.
while count <= 10000:
    try:
        num = int(input())
    except: break
    pre.append(num)
    count += 1

post_order(0, len(pre)-1)