[백준] 트리 4256번 Python

문제

트리
BT의 전위, 중위 순회한 결과가 주어졌을 때, 후위 순회했을 때의 결과를 구하는 프로그램을 작성하시오.

 

첫째 줄에 테스트 케이스의 개수 T가 주어진다. 각 테스트 케이스의 첫째 줄에는 노드의 개수 n이 주어진다. (1 ≤ n ≤ 1,000) BT의 모든 노드에는 1부터 n까지 서로 다른 번호가 매겨져 있다. 다음 줄에는 BT를 전위 순회한 결과, 그 다음 줄에는 중위 순회한 결과가 주어진다. 항상 두 순회 결과로 유일한 이진 트리가 만들어지는 경우만 입력으로 주어진다.

각 테스트 케이스마다 후위 순회한 결과를 출력 한다.

문제 생각 & 풀이

주어진 전위 순회(preorder)와 중위 순회(inorder) 결과를 통해 이진 트리를 복원하고, 후위 순회(postorder) 결과를 구하는 것이다. 재귀를 이용한 트리 복원 및 후위 순회 결과를 출력해주는 방식으로 시도하였다.

전위 순회의 첫 번째 요소는 항상 서브트리의 루트로 해주고, 중위 순회를 통해 루트 노드의 왼쪽 서브트리와 오른쪽 서브트리를 구분해주었다. 그리고 이를 기반으로 트리를 재귀적으로 복원하며 후위 순회를 수행하는 방식으로 시도하니 쉽게 풀 수 있었다 !

 

• 입력 부분:
  • 노드의 개수 n, 전위 순회 리스트 preorder, 중위 순회 리스트 inorder를 입력받음.
• 재귀 함수 정의 (solve):
  • 전위 순회의 첫 번째 요소는 항상 루트 노드다.
  • 중위 순회에서 루트 노드의 위치를 찾아 왼쪽 서브트리와 오른쪽 서브트리로 나누어서 진행
  • 각 서브트리에 대해 재귀적으로 solve 함수를 호출하여 후위 순회 결과를 구한다.
  • 후위 순회 결과는 왼쪽 서브트리, 오른쪽 서브트리, 루트 노드 순으로 결합해주어 return한다.

 

CODE

import sys
input = sys.stdin.readline

# 후위 순회 시작
def solve(preorder, inorder):
  if not preorder:
    return []

  root = preorder[0] # 전위 순회 처음 요소 = 루트 노드
  root_idx = inorder.index(root) # 중위 순회 루트 노드 인덱스

  # 루트 노드가 기준
  inorder_left = inorder[:root_idx]  # 왼쪽 서브트리
  inorder_right = inorder[root_idx+1:]  # 오른쪽 서브트리

  preorder_left = preorder[1:1+len(inorder_left)]  # 왼쪽 서브트리
  preorder_right = preorder[1+len(inorder_left):]  # 오른쪽 서브트리

  # 후위 순회 결과 = 왼쪽 서브트리 + 오른쪽 서브트리 + 루트 노드
  return solve(preorder_left, inorder_left) + solve(preorder_right, inorder_right) + [root]


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

  res = solve(preorder, inorder)
  print(" ".join(map(str, res)))