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

트리 문제는 언제 풀어도 어려운 것 같다. 전위순회 중위순회 후위순회 뭔 어려운 말들이 많은지...

 

그래도 지난번에 이진 검색 트리 문제를 풀었기에 그때의 기억을 살려서 풀 수 있었다.

 

가장 중요한 포인트는 인오더와 포스트오더 둘다 왼쪽 트리를 먼저 순회한다는 점, 그리고 포스트오더의 가장 마지막 지점이 부모노드라는 점, 그리고 인오더는 부모노드를 기준으로 왼쪽트리와 오른쪽트리로 나누어진다는 점이다.

 

따라서 포스트오더의 마지막 값인 부모노드가 인오더에서 위치한 지점을 찾고, 그 점을 기준으로 리스트를 둘로 나누어 재귀함수를 돌아주면 된다.

 

애먹었던 부분은 리스트를 나누어 주다가 리스트의 길이가 0 또는 1이 될 때 어떻게 처리해주어야 할지였는데, 0일때는 return, 1일때는 값을 추가하고 return을 하니 해결되었다.

 

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

N = int(input())

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

preorder = deque()

def makepre(start1,end1,start2,end2):
  if start1>end1:
    return
  preorder.append(postorder[end2])
  if start1==end1:
    return
  s = inorder.index(postorder[end2])-start1
  makepre(start1,s+start1-1,start2,s+start2-1)
  makepre(s+start1+1,end1,s+start2,end2-1)

makepre(0,N-1,0,N-1)
print(" ".join(map(str,preorder)))

 

위 코드는 python3에서는 통과되지 못하고 pypy3에서만 통과한 코드이다. 3152ms가 걸렸다.

시간이 소모되는 지점은 inorder.index 함수 때문이었다. index함수의 시간복잡도가 O(N)이기 때문.

 

이를 어떻게 처리할까 고민하다가 위치에 대한 dictionary를 처음부터 만들어주면 된다는 것을 깨달았다.

 

아래는 수정한 코드이다.

import sys
input = sys.stdin.readline
from collections import deque
sys.setrecursionlimit(10**6)

N = int(input())

inorder = list(map(int,input().split()))
postorder = list(map(int,input().split()))
Index = {}
for i in range(N):
  Index[inorder[i]] = i

preorder = deque()

def makepre(start1,end1,start2,end2):
  if start1>end1:
    return
  preorder.append(postorder[end2])
  if start1==end1:
    return
  s = Index[postorder[end2]]-start1
  makepre(start1,s+start1-1,start2,s+start2-1)
  makepre(s+start1+1,end1,s+start2,end2-1)

makepre(0,N-1,0,N-1)
print(" ".join(map(str,preorder)))

Index라는 dictionary를 만들어 각 노드의 위치를 저장하고 index함수 대신 실행한다.

이렇게 바꾸니 시간이 356ms로 10분의 1로 단축되었다.

 

 

오늘의 교훈) 항상 시간복잡도를 해결할 방법을 모색하자