[재귀/파이썬] 백준 4256번: 트리 / 골드 3

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

4256번: 트리

첫째 줄에 테스트 케이스의 개수 T가 주어진다. 각 테스트 케이스의 첫째 줄에는 노드의 개수 n이 주어진다. (1 ≤ n ≤ 1,000) BT의 모든 노드에는 1부터 n까지 서로 다른 번호가 매겨져 있다. 다음

www.acmicpc.net

<최초 코드>

import sys
input = sys.stdin.readline

class BT():
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def recursion(preorder, inorder):
    if not preorder:
        return None

    bt = BT()
    node = preorder[0]
    bt.val = node

    idx = inorder.index(node)
    left = preorder[1:idx+1]
    right = preorder[idx+1:]

    bt.left = recursion(left, inorder[:idx])
    bt.right = recursion(right, inorder[idx+1:])
    return bt

def postorder(bt):
    if not bt:
        return

    postorder(bt.left)
    postorder(bt.right)
    post.append(bt.val)

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

    bt = recursion(preorder, inorder)

    post = []
    postorder(bt)
    print(*post)

굳이 트리 자료구조까지 만들고, 다시 후위순회 코드를 만들어야될까 싶어서 조금 더 생각을 해봤다.

최초 코드를 조금 더 다듬어 코드를 확 줄일 수 있었다.

<개선 코드>

import sys
input = sys.stdin.readline

def recursion(preorder, inorder):
    if not preorder:
        return

    node = preorder[0]
    idx = inorder.index(node)

    recursion(preorder[1:idx+1], inorder[:idx])
    recursion(preorder[idx+1:], inorder[idx+1:])
    print(node, end=' ')

t = int(input())
for _ in range(t):
    n = int(input())
    preorder = list(map(int, input().split()))
    inorder = list(map(int, input().split()))
    recursion(preorder, inorder)
    print()