2263. 트리의 순회 (Python)

https://www.acmicpc.net/problem/2263

2263번: 트리의 순회

 

풀이

 입력으로 주어진 인오더와 포스트오더를 사용하여 재귀를 통해 프리오더를 구할 수 있다. 포스트오더는 리스트의 마지막에 있는 숫자가 현재 트리의 부모 노드에 해당한다. 이를 통해서 현재 트리의 부모 노드 숫자를 알 수 있다. 이제 이 숫자를 인오더 리스트의 어느 위치에 있는지를 찾는다. 인오더의 위치를 기준으로 왼쪽에 있는 숫자 리스트가 현재 트리의 왼쪽 서브 트리에 해당하고, 오른쪽에 있는 숫자 리스트가 현재 트리의 오른쪽 서브 트리에 해당한다.

 

 이를 프리오더에 맞게 먼저 루트에 해당하는 노드의 값을 출력한다. 그 후 왼쪽 서브트리와 오른쪽 서브트리에 대하여 차례로 재귀 호출하여 똑같이 부모 노드를 찾고 다시 왼쪽 서브트리, 오른쪽 서브트리에 대하여 재귀 호출을 반복하면 프리오더를 구할 수 있다.

 

 문제를 풀 때 생각해 봐야 할 것들이 있다. 리스트에 대해서 슬라이싱을 사용하면 메모리를 잡아먹게 되고 해당 문제에서는 메모리 초과가 발생하게 된다. 그렇기 때문에 슬라이싱이 아니라 인덱스 값을 직접 구해서 함수의 인자 값으로 전달하는 방식을 사용하자.

 

 그리고 포스트오더에서 찾은 부모 노드의 값을 인오더에서 찾을 때 다음과 같이 index 함수를 사용할 수 있다.

# 부모 노드에 해당하는 값이 inOrder 리스트의 몇 번재 인덱스에 있는지 찾는다.
rootIndex = inOrder.index(postOrder[-1])

 하지만 이럴 경우 많은 시간이 걸린다. 문제에서 값이 1부터 n까지 중복 없이 주어진다고 했으므로 각 값이 inOrder에서 몇 번째 인덱스에 있는지를 따로 리스트를 정의하여 구해놓고 이를 사용하면 시간을 크게 단축시킬 수 있다.

 

코드

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


# inOrder 리스트의 인덱스 시작과 끝,
# postOrder 리스트의 인덱스 시작과 끝을 함수의 파라미터로 전달한다. 
def preOrder(inorderStart, inorderEnd, postStart, postEnd):
    if inorderStart > inorderEnd:
        return

    if inorderStart == inorderEnd:
        ans.append(inOrder[inorderEnd]) # 부모 노드의 값을 ans 리스트에 추가 
        return

    rootIdx = inorderIndex[postOrder[postEnd]] # 부모 노드의 값이 inOrder 리스트의 몇 번째에 있는지를 구한다.
    inorderLeftLength = rootIdx - 1 - inorderStart # 왼쪽 서브트리의 길이 

    ans.append(postOrder[postEnd]) # 부모 노드의 값을 ans 리스트에 추가 
    preOrder(inorderStart, rootIdx - 1, postStart, postStart + inorderLeftLength) # 왼쪽 서브트리
    preOrder(rootIdx + 1, inorderEnd, postStart + inorderLeftLength + 1, postEnd - 1) # 오른쪽 서브트리


if __name__ == '__main__':
    N = int(input())
    inOrder = list(map(int, input().split())) # 인오더
    postOrder = list(map(int, input().split())) # 포스트오더
    ans = []  # 프리오더
    inorderIndex = [0 for _ in range(N+1)] # 각 값이 inOrder 리스트의 몇 번째에 있는지에 대한 리스트 

    # inorderIndex[n] = m -> 숫자 n이 inOrder[m]에 있음을 의미 
    # 시간을 단축시키기 위한 작업 
    for i in range(N):
        inorderIndex[inOrder[i]] = i 

    preOrder(0, N - 1, 0, N - 1)
    print(*ans)