[백준/4256] 트리 (Python)

Problem : https://www.acmicpc.net/problem/4256

 

Difficulty : Gold 2

 

Status : Solved

 

Time : 00:13:02

 

---

 

문제 설명

 

이진 트리는 매우 중요한 기본 자료 구조이다. 아래 그림은 루트 노드가 유일한 이진 트리이다. 모든 노드는 최대 2개의 자식 노드를 가질 수 있으며, 왼쪽 자식이 순서가 먼저이다. 노드 n개로 이루어진 이진 트리를 BT라고 하자. BT의 노드는 1부터 n까지 유일한 번호가 매겨져 있다.

아래 그림에 나와있는 BT의 루트는 3번 노드이다. 1번 노드는 오른쪽 자식만 가지고 있고, 4와 7은 왼쪽 자식만 가지고 있다. 3과 6은 왼쪽과 오른쪽 자식을 모두 가지고 있다. 나머지 노드는 모두 자식이 없으며, 이러한 노드는 리프 노드라고 부른다.

 

BT의 모든 노드를 순회하는 방법은 전위 순회(preorder), 중위 순회(inorder), 후위 순회(postorder)로 총 세 가지가 있다. 이 세 방법은 아래에 C 스타일의 의사 코드로 나와 있다. BT의 노드 v에 대해서, v.left는 왼쪽 자식, v.right는 오른쪽 자식을 나타낸다. v가 왼쪽 자식이 없으면 v.left는 ∅와 같고, 오른쪽 자식이 없으면 v.right는 ∅와 같다.

BT를 전위 순회, 중위 순회한 결과가 주어진다. 즉, 위의 함수 중 preorder(root node of BT)와 inorder(root node of BT)를 호출해서 만든 리스트가 주어진다. 두 순회한 결과를 가지고 다시 BT를 만들 수 있다. BT의 전위, 중위 순회한 결과가 주어졌을 때, 후위 순회했을 때의 결과를 구하는 프로그램을 작성하시오.

예를 들어, 위의 그림을 전위 순회하면 3,6,5,4,8,7,1,2, 중위 순회하면 5,6,8,4,3,1,2,7이 된다. 이를 이용해 후위 순회하면 5,8,4,6,2,1,7,3이 된다.

 

입력 및 출력

 

입력

 

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

 

출력

 

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

 

입력 예시

 

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

 

출력 예시

 

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

 

 

---

 

풀이

 

전위 순회와 중위 순회의 특징을 알면 쉽게 접근해볼 수 있는 문제이다.

 

트리의 서브트리에 대해 차례대로 inorder, preoder, postorder 순회 결과를 위와 같이 나타낼 수 있다. 여기서 우리의 목적은 preoder와 inorder의 정보만을 가지고 postorder순으로 나타내는 것이다.

 

여기서 잠깐. 다음과 같은 특징들을 발견해 볼 수 있을 것이다.

 

• preorder의 가장 첫 번째 노드는 무조건 루트 노드이다. 루트 노드를 제일 먼저 순회하기 때문이다!
• 그렇다면 inorder상에서 root의 위치를 찾아볼 수도 있을 것이다. 여기서 두 번째. inorder의 root 왼쪽에 속하는 노드는 무조건 left subtree, 오른쪽에 속하는 노드는 무조건 right subtree이다!
• 따라서 위 두 과정을 순서대로 거치면, 우리는 root, left subtree에 속하는 노드들, right subtree에 속하는 노드들을 분리해낼 수 있다. 즉 이 left subtree, right subtree에 대해 재귀적으로 postorder 순으로 변환하고, left-right-root순으로 나열하여 반환한다면 우리는 결과적으로 이 서브트리의 postorder를 구할 수 있다!

 

풀이 코드

import sys
input = sys.stdin.readline

def make_postorder(preorder, inorder) :
  if not preorder :
    return list()
  root = preorder[0]
  idx = inorder.index(root)
  left = make_postorder(preorder[1:idx+1], inorder[:idx])
  right = make_postorder(preorder[idx+1:], inorder[idx+1:])
  return left + right + [root]

for _ in range(int(input())) :
  n = int(input())
  preorder = list(map(int, input().split()))
  inorder = list(map(int, input().split()))
  result = make_postorder(preorder, inorder)
  print(*result)

풀이 완료~

리스트 슬라이싱을 가지고 풀이하였는데, 사실 리스트 인덱싱을 통해서도 충분히 풀어볼 수 있는 문제이다. 이 방법을 사용하면 좀더 시간을 단축해 볼 수 있을 것으로 기대된다!