[백준 2250번] 트리의 높이와 너비 (python)

문제

이진트리를 다음의 규칙에 따라 행과 열에 번호가 붙어있는 격자 모양의 틀 속에 그리려고 한다. 이때 다음의 규칙에 따라 그리려고 한다.

1. 이진트리에서 같은 레벨(level)에 있는 노드는 같은 행에 위치한다.
2. 한 열에는 한 노드만 존재한다.
3. 임의의 노드의 왼쪽 부트리(left subtree)에 있는 노드들은 해당 노드보다 왼쪽의 열에 위치하고, 오른쪽 부트리(right subtree)에 있는 노드들은 해당 노드보다 오른쪽의 열에 위치한다.
4. 노드가 배치된 가장 왼쪽 열과 오른쪽 열 사이엔 아무 노드도 없이 비어있는 열은 없다.

이와 같은 규칙에 따라 이진트리를 그릴 때 각 레벨의 너비는 그 레벨에 할당된 노드 중 가장 오른쪽에 위치한 노드의 열 번호에서 가장 왼쪽에 위치한 노드의 열 번호를 뺀 값 더하기 1로 정의한다. 트리의 레벨은 가장 위쪽에 있는 루트 노드가 1이고 아래로 1씩 증가한다.

아래 그림은 어떤 이진트리를 위의 규칙에 따라 그려 본 것이다. 첫 번째 레벨의 너비는 1, 두 번째 레벨의 너비는 13, 3번째, 4번째 레벨의 너비는 각각 18이고, 5번째 레벨의 너비는 13이며, 그리고 6번째 레벨의 너비는 12이다.

우리는 주어진 이진트리를 위의 규칙에 따라 그릴 때에 너비가 가장 넓은 레벨과 그 레벨의 너비를 계산하려고 한다. 위의 그림의 예에서 너비가 가장 넓은 레벨은 3번째와 4번째로 그 너비는 18이다. 너비가 가장 넓은 레벨이 두 개 이상 있을 때는 번호가 작은 레벨을 답으로 한다. 그러므로 이 예에 대한 답은 레벨은 3이고, 너비는 18이다.

임의의 이진트리가 입력으로 주어질 때 너비가 가장 넓은 레벨과 그 레벨의 너비를 출력하는 프로그램을 작성하시오.

 

기록

(트리 5일차) 벌써 트리 5일차.. 인데도 막히는 부분이 많은 것 같다. 계속 제출하자마자 틀렸습니다 뜨고 여러 테스트케이스 넣어봐도 안되길래 찾아봤더니 루트가 1이 아니라는 것을 알게됐다.....! '트리의 레벨은 가장 위쪽에 있는 루트 노드가 1이고 아래로 1씩 증가한다.' 문제 중 이 부분 읽으면서 당연히 루트노드가 1인줄 알았는데 다시 읽어보니까 루트 노드의 레벨이 1이라는 말이었다. 그래서 루트노드 찾는 코드 구현하고 성공했다:)

 

내가 생각한 이 문제의 키포인트는 (루트노드가 1이 아니라는 것이랑) 트리 각 노드의 열번호가 inorder로 순회하는 순서라는 점인 것 같다!

 

코드

from sys import stdin
from collections import defaultdict

def make_coldict_by_level(tree, root, level):

    # base case
    if root==-1 : return
    
    # step
    global coldict, col
    make_coldict_by_level(tree, tree[root][0], level+1)
    coldict[level].append(col)
    col+=1
    make_coldict_by_level(tree, tree[root][1], level+1)
    return


# 변수 초기화
N = int(stdin.readline())
tree = {}
tree_up = {}
for i in range(N):
    n, l, r = map(int, stdin.readline().split())
    tree[n] = (l, r)
    tree_up[r] = n
    tree_up[l] = n

# root 찾기
root = -1
while root in tree_up:
    root = tree_up[root]

# 각 레벨마다 노드 열번호 배열이 저장된 dict 생성
coldict = defaultdict(list)
col = 1
make_coldict_by_level(tree, root, 1)

# dict 돌면서 max width 찾기
maxwidth = 0
maxlevel = 1
for level, cols in sorted(coldict.items()):
    width = cols[-1] - cols[0]
    if width>maxwidth:
        maxwidth = width
        maxlevel = level

print(maxlevel, maxwidth+1)

1. 먼저 입력값들 변수에 저장한다. 루트-자식노드들로 이루어진 기본 트리 dict, 루트 찾기위해 자식노드-루트로 이루어진 dict 이렇게 두 개 생성했다.

 

2. 루트 찾기 수행

 

3. 각 레벨마다 노드들의 열번호가 담긴 dict 생성 {level : [node1_col, node2_col], ,}

make_coldict_by_level은 tree를 탐색하는 재귀함수로, 현재 레벨을 key로 하여 각 노드들의 열번호(preorder 탐색순서)를 배열로 저장한다. 현재 preorder 순서번호와 dict를 전역변수 안쓰고 해보려고했는데 복잡해져서 편하게 전역변수 선언했다ㅎ

 

4. 최대 너비 찾기

level 순서로 dict 돌면서 마지막노드의 열번호 - 첫번째노드의 열번호로 너비를 계산하여 최댓값을 찾는다.