[백준] 2250번 트리의 높이와 너비 - Python

이번에 해결해 볼 문제는 백준에 있는 2250번 트리의 높이와 너비입니다.

---

문제 설명

1. 이진트리를 다음의 규칙에 따라 행과 열에 번호가 붙어있는 격자 모양의 틀 속에 그리고자 합니다.
   • 이진트리에서 같은 레벨(level)에 있는 노드는 같은 행에 위치합니다.
   • 한 열에는 한 노드만 존재합니다.
   • 임의의 노드의 왼쪽 부트리(left subtree)에 있는 노드들은 해당 노드보다 왼쪽의 열에 위치하고, 오른쪽 부트리(right subtree)에 있는 노드들은 해당 노드보다 오른쪽의 열에 위치합니다.
   • 노드가 배치된 가장 왼쪽 열과 오른쪽 열 사이엔 아무 노드도 없이 비어있는 열은 없습니다.
2. 가장 너비가 큰 레벨을 구합니다.
   • 너비 공식 = 그 레벨에 할당된 노드 중 가장 오른쪽에 위치한 노드의 열 번호에서 가장 왼쪽에 위치한 노드의 열 번호를 뺀 값 더하기 1
3. 너비가 가장 넓은 레벨과 그 너비를 차례로 출력하면 됩니다.

문제 해결 아이디어

• 구해야 할 것을 정리하면 아래와 같습니다.
  • 각 node의 레벨
  • 각 node의 column index
• 먼저 각 node의 column index는 중위 순회를 통해 구할 수 있습니다.
• 각 node의 레벨은 중위 순회를 돌면서 재귀함수의 특징을 사용해 구할 수 있습니다.

주의해야 할 점

• 루트 노드의 번호는 정해져 있지 않습니다.
• 각 노드의 레벨이 순차적이지 않습니다.
  • 1번 노드의 레벨은 4, 2번 노드의 레벨은 1일 수 있습니다.

문제 해결 코드

import sys
from collections import defaultdict

input = sys.stdin.readline

N = int(input())

tree = defaultdict(list)
parents = [False] * (N + 1)
for _ in range(N):
    node, left, right = map(int, input().split())
    tree[node].append(left)
    tree[node].append(right)
    if left != -1:
        parents[left] = True
    if right != -1:
        parents[right] = True

# 루트 노드 구하기
root = 0
for i in range(1, N + 1):
    if not parents[i]:
        root = i
        break

level = [0] * (N + 1)
visited = [0]


def inorder(node, lev_val):
    if node != -1:
        lev_val += 1
        level[node] = lev_val # level 구하기

        inorder(tree[node][0], lev_val)
        visited.append(node) # column index 구하기
        inorder(tree[node][1], lev_val)


inorder(root, 0)

col = [0] * (N + 1)

for i in range(1, N + 1):
    col[visited[i]] = i


# 각 레벨별로 node를 탐색하면서 너비 구하기
answer_val = 0
answer_lev = 0
for i in range(1, max(level) + 1):
    min_col = N + 1
    max_col = 0
    for j in range(1, N + 1):
        if level[j] == i:
            min_col = min(min_col, col[j])
            max_col = max(max_col, col[j])
    width = max_col - min_col + 1
    if width > answer_val:
        answer_val = width
        answer_lev = i
print(answer_lev, answer_val)