[Python] 백준 2263번: 트리의 순회

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

2263번: 트리의 순회

 

import sys
sys.setrecursionlimit(500000)
input = sys.stdin.readline


def printPreorder(inorderStartIndex, inorderEndIndex, postorderStartIndex, postorderEndIndex):
    length = inorderEndIndex - inorderStartIndex

    if length == 0:
        return
    elif length == 1:
        print(inorder[inorderStartIndex:inorderEndIndex][0], end=" ")
        return

    root = postorder[postorderEndIndex-1]
    inorderRootIndex = inorderIndices[root]
    leftLength = inorderRootIndex - inorderStartIndex

    print(root, end=" ")
    printPreorder(inorderStartIndex, inorderRootIndex,
                  postorderStartIndex, postorderStartIndex + leftLength)
    printPreorder(inorderRootIndex + 1, inorderEndIndex,
                  postorderStartIndex + leftLength, postorderEndIndex - 1)


if __name__ == '__main__':
    # 입력
    n = int(input())
    inorder = list(map(int, input().split()))
    postorder = list(map(int, input().split()))

    # 각 노드의 inorder에서의 인덱스를 저장
    inorderIndices = [-1]*(n+1)
    for i in range(n):
        inorderIndices[inorder[i]] = i

    # 출력
    printPreorder(0, n, 0, n)

 

아래 예시를 기준으로 설명한다.

8
4 2 1 7 8 5 3 6
4 2 8 7 5 6 3 1

답은 아래와 같다.

우선 특징을 분석해보자.

- postorder의 마지막 원소는 무조건 루트이다.

- inorder에서는 루트를 기준으로 양옆을 가르면 left, right 요소가 나뉜다.

- inorder을 통해 left, right의 길이를 얻으면, 해당 길이를 이용하여 postorder을 앞에서부터 slice하여 left, right를 얻을 수 있다.

 

해당 특징들을 토대로 코드의 작성 방향을 잡아보자. 우선 전체 범위에 대해서 root를 찾았고, left와 right을 나누었다고 하자. 그럼 다음에는 또 다시 left와 right에 대해서 같은 행위를 반복해야한다. 즉, 재귀 함수를 사용할 것이다.

또한 inorder과 postorder의 범위를 계속해서 작은 범위로 나누어가면서 진행할 것이다. 재귀함수를 통한 분할정복을 수행하겠다는 의미이다. 가령, 위에서 left, right가 (4, 2)와 (7, 8, 5, 3, 6)이 나누어졌고, 여기서 재귀함수를 각각 (4, 2)와 (7, 8, 5, 3, 6)의 범위에서 또다시 진행한다는 것이다. 여기서 선택지가 2개가 있다.

1. inorder, postorder이라는 배열을 계속해서 slice해가면서 재귀함수를 수행함

2. inorder, postorder은 그대로 두고, 각각의 범위를 나타내는 정수를 전달하는 방식으로 수행함

여기서 내 풀이는 2번에 해당하는데, 1번의 경우 메모리 초과가 나기 때문이다. 배열을 slice하면 해당하는 새로운 배열이 할당되기 때문이다.

 

다음으로, 재귀 함수 내부(printPreorder)를 구체적으로 설계해보자. 베이스 케이스는 제쳐두고, 재귀 케이스부터 고려한다.

1. 해당 범위 내에서의 root 노드를 찾는다.

2. inorder에서 root의 인덱스를 찾는다. 다음 3번의 작업에서, 해당 인덱스를 기준으로 inorder의 좌우를 나누고, 또한 left의 총길이(left가 4, 2이면 총길이는 2)를 나타내는 leftLength를 구할 수 있다.

3. preorder 순서로 출력한다. (root -> printPreorder를 left에 대해 수행 -> printPreorder를 right에 대해 수행)

3번에서 범위를 나누는 것만 정상적으로 작성한다면 크게 어려울 것 없다.

 

마지막으로 베이스 케이스이다. 우선 현재 함수에서 수행해야할 배열의 길이를 알아야한다. 함수의 인자로 들어오는 범위의 차이가 곧 현재 함수에 해당하는 배열의 길이(length)를 뜻한다. 어차피 재귀 함수에 전달되는 inorder과 postorder의 크기는 같으므로, length에 (inorderEndIndex - inorderStartIndex)가 오든, (postorderEndIndex - postorderStartIndex)가 오든 상관이 없다.

먼저 길이가 0이면 출력하지 않고 바로 리턴해준다.

길이가 1이면 해당 노드를 곧바로 출력해주고 리턴해준다.