[백준] 9934 완전 이진 트리

문제 설명

문제 링크

해당 문제는 전위 순회로 탐색한 트리의 노드 방문 순서가 입력으로 주어지며, 이를 토대로 트리의 각 레벨을 위에서부터 아래로 그리는 문제이다.

문제 풀이

처음에는 단순하게 중위 순회로 주어진 것을 전위 순회로 바꾼 후 출력하면 된다고 생각했다. 그렇게 하면 루트 노드가 가장 먼저 나오고 레벨대로 나올 거라고 생각했다.

import sys
from collections import deque
input = sys.stdin.readline

# L D R 중위순회

K = int(input())
inorder = list(map(int, input().split()))
preorder = deque()
#를 D L R 전위순회로 바꾸어라 

def change_to_preorder(inorder, preorder):
    if not inorder:
        return
    mid = len(inorder)//2
    preorder.append(inorder[mid])
    change_to_preorder(inorder[:mid], preorder)
    change_to_preorder(inorder[mid+1:], preorder)

change_to_preorder(inorder, preorder)

for i in range(K):
    for j in range(2**i):
        if preorder:
            print(preorder.popleft(), end = "")
            if j != (2**i - 1):
                print(" ", end ="")
    print()

예제 1은 맞게 나왔는데, 계속 틀렸다고 떠서 의문이던 찰나에, 혹시나 해서 예제 2를 넣어 보니까 각 레벨을 순서대로 출력하는 게 아니라 전위 순회로 출력하여 다음과 같이 나왔다:

3
6 1
4 2 5 7 

그렇다. 뻘짓을 한것이었다. 물론 전위 순회로 바꾸었기 때문에 이는 DFS를 했을 떄와 같을테고, 이 DFS를 BFS로 바꾸기만 하면 사실 각 레벨을 표현할 수 있을 것이라 생각했지만, 일단 처음부터 다시 해보기로 했다.

각 층을 분리하기 위해 카운터가 필요했으며, 층별 별도 공간을 만들기 위해 이차원 배열을 사용했다.

# 레벨에 해당하는 배열 생성
levels = [[] for _ in range(K)]

def make_level(arr, current_level, levels):

기존에 전위 순회로 바꾸던 코드에서는 현재 층을 표현하는 current_level이 없었기에 재귀를 타도 D-L-R 순서로 좌측 서브트리를 먼저 전체 순회한 후에야 우측 서브트리를 순회했기 때문에 각 서브트리의 루트노드들을 먼저 수집할 방법이 없었다.

위와 같은 형식으로 current_level 카운터로 층을 표현하고, 각 층을 2차원 배열 안에 저장하면 재귀를 탈 때마다 current_level + 1을 인자로 주면 층별 노드를 관리할 수 있다.

코드

import sys
input = sys.stdin.readline

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

# 레벨에 해당하는 배열 생성
levels = [[] for _ in range(K)]
 
def make_level(arr, current_level, levels):
    if not arr:
        return

    mid = len(arr)//2 

    levels[current_level].append(arr[mid])

    make_level(arr[:mid], current_level + 1, levels)
    make_level(arr[mid+1:], current_level + 1, levels)

make_level(inorder, 0, levels)

for level in levels:
    print(*level, sep= " ")