[백준 4256번] 트리 (python)

문제

이진 트리는 매우 중요한 기본 자료 구조이다. 아래 그림은 루트 노드가 유일한 이진 트리이다. 모든 노드는 최대 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이 된다.

 

기록

(트리 4일차) 계속 비슷한 문제들 풀다보니까 트리문제 조금씩 감이 잡히는 것 같다. 이 문제 키포인트는 preorder일 때 가장 첫번째 요소가 루트 노드인 것, inorder일 때 루트노드 기준으로 왼쪽이 왼쪽자식노드, 오른쪽이 오른쪽자식노드라는 것!

그리고 문제 풀다가 갑자기 시간제한 1초인데 넘어가도 맞게 나와서 궁금해서 찾아봤다. 언어별로 주어지는 시간이 다르구나..

 

나중에 코드 좀 개선했더니 시간이 많이 줄었다! (코드는 아래에)

 

코드

from sys import stdin

def get_postlist(prelist, pstart, pend, inlist, istart, iend):

    # Base case (자식 노드 없으면 현재 노드 반환)
    if pend - pstart<=1 : return prelist[pstart:pend]

    # Step (루트 / 왼쪽자식 / 오른쪽자식 구분한 뒤 후위순서로 조합해서 반환)
    # 1. preorder 맨 앞이 루트노드
    root = prelist[pstart]
    # 2. inorder에서 루트 기준으로 왼쪽/ 오른쪽 자식노드 구분
    iroot = 0
    for i in range(istart, iend):
        if inlist[i]==root: iroot = i
    # 3. 위에서 찾은 왼쪽/오른쪽 자식노드 확인하면서 preorder도 자식노드 구분 
    pmid = pstart+1
    for i in range(pstart+1, pend):
        if prelist[i] not in inlist[istart:iroot]:
            break
        pmid += 1
    # 4. 구분한 루트/왼쪽/오른쪽 노드 후위순서로 조합
    postlist = []
    postlist += get_postlist(prelist, pstart+1, pmid, inlist, istart, iroot)
    postlist += get_postlist(prelist, pmid, pend, inlist, iroot+1, iend)
    postlist += [root]
    return postlist


T = int(stdin.readline())
for i in range(T):

    # 변수 초기화
    N = int(stdin.readline())
    prelist = list(map(int, stdin.readline().split()))
    inlist = list(map(int, stdin.readline().split()))

    # 후위변환한 결과 출력
    postlist = get_postlist(prelist, 0, len(prelist), inlist, 0, len(inlist))
    print(' '.join(map(str, postlist)))

먼저 preorder 입력값은 prelist에, inorder 입력값은 inlist에 저장했다. 

그리고 각 노드마다 서브트리의 왼쪽자식 / 오른쪽자식 / 루트 구분해서 postorder로 조합해서 반환하는 재귀함수 구현한 뒤, 그 결과 출력하도록 구현했다.

 

*수정

def get_postlist(prelist, pstart, pend, inlist, istart, iend):

    # Base case (자식 노드 없으면 현재 노드 반환)
    if pend - pstart<=1 : return prelist[pstart:pend]

    # Step (루트 / 왼쪽자식 / 오른쪽자식 구분한 뒤 후위순서로 조합해서 반환)
    # 1. preorder 맨 앞이 루트노드
    root = prelist[pstart]
    # 2. inorder에서 루트노드 인덱스 찾기
    iroot = inlist.index(root)
    # 3. preorder도 자식노드 구분하기 위해 중간 인덱스 찾기
    pmid = (pstart+1)+(iroot-istart)
    # 4. 구분한 루트/왼쪽/오른쪽 노드 후위순서로 조합
    postlist = []
    postlist += get_postlist(prelist, pstart+1, pmid, inlist, istart, iroot)
    postlist += get_postlist(prelist, pmid, pend, inlist, iroot+1, iend)
    postlist += [root]
    return postlist

후위순서로 변경한 리스트 반환하는 재귀함수 코드 개선했다.

Step 2번의 inorder에서 루트노드 인덱스 찾기를 함수써서 코드 줄였고,

Step 3번의 preorder 왼쪽/오른쪽 자식노드 구분하는 인덱스 찾는 과정을 루트노드+1 인덱스부터(pstart+1) 왼쪽자식노드 개수(iroot-istart) 더하는 수식으로 변경했다.

 

이렇게 수정했더니 시간이 2064ms -> 792ms로 줄었다!