[백준] (2263) 트리의 순회 [Python]

문제 링크

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

문제

n개의 정점을 갖는 이진 트리의 정점에 1부터 n까지의 번호가 중복 없이 매겨져 있다. 이와 같은 이진 트리의 인오더와 포스트오더가 주어졌을 때, 프리오더를 구하는 프로그램을 작성하시오.

• 입력 : 첫째 줄에 n(1 ≤ n ≤ 100,000)이 주어진다. 다음 줄에는 인오더를 나타내는 n개의 자연수가 주어지고, 그 다음 줄에는 같은 식으로 포스트오더가 주어진다.
• 출력 : 첫째 줄에 프리오더를 출력한다.

예제 입력	예제 출력
11 8 4 9 2 5 1 10 6 11 3 7 8 9 4 5 2 10 11 6 7 3 1	1 2 4 8 9 5 3 6 10 11 7

---

코드

(1) 메모리 초과

n = int(input())
inorder = list(map(int, input().split()))
postorder = list(map(int, input().split()))
preorder = list()

def solution(arr_i, arr_p):
    if not arr_i or not arr_p:
        return
    root = arr_p[-1]
    mid = arr_i.index(root)
    
    left_i = arr_i[:mid]
    right_i = arr_i[mid+1:]
    left_p = arr_p[:len(left_i)]
    right_p = arr_p[len(left_i):len(left_i)+len(right_i)]
    
    preorder.append(root)
    solution(left_i, left_p)
    solution(right_i, right_p)

solution(inorder, postorder)
print(*preorder)

 

(2) 정답 코드

import sys
sys.setrecursionlimit(10**6)

n = int(input())
inorder = list(map(int, input().split()))
postorder = list(map(int, input().split()))
preorder = list()

pos = [0] * (n+1)
for i in range(n):
    pos[inorder[i]] = i

def solution(in_l, in_r, post_l, post_r):
    if in_l > in_r or post_l > post_r:
        return
    root = postorder[post_r]
    mid = pos[root]
    left_size = mid - in_l
    
    preorder.append(root)
    solution(in_l, mid-1, post_l, post_l+left_size-1)
    solution(mid+1, in_r, post_l+left_size, post_r-1)

solution(0, n-1, 0, n-1)
print(*preorder)

 

재귀 호출마다 새로운 리스트 슬라이싱(arr_i[:mid], arr_i[mid+1:])을 생성하기 때문에 메모리 초과가 났다고 생각해서, 리스트를 매번 슬라이싱해서 주는 대신에 인덱스만 따로 저장해서 넘겨주는 방식으로 수정했다.

---

번외

비슷한 방식으로 프리오더와 인오더가 주어질 때 포스트오더를 구하는 문제(BOJ4256)는 다음과 같이 풀 수 있다.

import sys
sys.setrecursionlimit(10**6)

T = int(input())
for _ in range(T):
    n = int(input())
    preorder = list(map(int, input().split()))
    inorder = list(map(int, input().split()))
    postorder = list()

    pos = [0] * (n+1)
    for i in range(n):
        pos[inorder[i]] = i

    def solution(pre_l, pre_r, in_l, in_r):
        if pre_l > pre_r or in_l > in_r:
            return
        root = preorder[pre_l]
        mid = pos[root]
        left_size = mid - in_l
        
        solution(pre_l+1, pre_l+left_size, in_l, mid-1)
        solution(pre_l+left_size+1, pre_r, mid+1, in_r)
        postorder.append(root)

    solution(0, n-1, 0, n-1)
    print(*postorder)

• 루트 노드 위치: 후위 순회에서는 마지막 원소(postorder[post_r])가 루트였지만, 전위 순회에서는 첫 번째 원소(preorder[pre_l])가 루트다.
• 재귀 호출 순서: 전위 순회 배열에서 왼쪽 서브트리는 pre_l+1부터 pre_l+left_size까지, 오른쪽 서브트리는 pre_l+left_size+1부터 pre_r까지다.
• 출력 순서: 후위 순회는 왼쪽 → 오른쪽 → 루트 순서이므로, 재귀 호출을 모두 마친 후 answer.append(root)를 수행한다.