백준 2250 [트리의 높이와 너비] 파이썬

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

 

# 백준 2250 트리의 높이와 너비
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)

N = int(input())  # 노드 수 입력

# 상수로 LEFT, RIGHT 인덱스 명시
LEFT = 0
RIGHT = 1

# conn[i] = [left child, right child]
conn = [[-1, -1] for _ in range(N+1)]

# 각 노드의 서브트리 노드 수 저장 (좌/우)
child_nums = [[0, 0] for _ in range(N+1)]

# 각 노드의 부모 노드를 저장 (-1이면 루트)
parents = [-1 for _ in range(N+1)]

# 트리 구조 입력
for _ in range(1, N+1):
    parent, left_child, right_child = map(int, input().split())
    conn[parent][LEFT] = left_child
    conn[parent][RIGHT] = right_child

    # 부모 설정 (루트 찾기용)
    if left_child != -1:
        parents[left_child] = parent
    if right_child != -1:
        parents[right_child] = parent

# 루트 노드 찾기 (부모가 없는 노드)
def find_root(node):
    if parents[node] != -1:
        return find_root(parents[node])
    else:
        return node

# 각 노드의 좌/우 서브트리에 포함된 노드 수 계산 (중위 순회 전 위치 계산용)
def dfs(node):
    left_child = conn[node][LEFT]
    right_child = conn[node][RIGHT]

    # 양쪽 자식이 있는 경우
    if left_child != -1 and right_child != -1:
        dfs(left_child)
        child_nums[node][LEFT] = 1 + child_nums[left_child][LEFT] + child_nums[left_child][RIGHT]

        dfs(right_child)
        child_nums[node][RIGHT] = 1 + child_nums[right_child][LEFT] + child_nums[right_child][RIGHT]

    # 왼쪽 자식만 있는 경우
    elif left_child != -1:
        dfs(left_child)
        child_nums[node][LEFT] = 1 + child_nums[left_child][LEFT] + child_nums[left_child][RIGHT]

    # 오른쪽 자식만 있는 경우
    elif right_child != -1:
        dfs(right_child)
        child_nums[node][RIGHT] = 1 + child_nums[right_child][LEFT] + child_nums[right_child][RIGHT]

    # 리프 노드일 경우 아무것도 하지 않음 (초기값이 0)

# 루트 노드에서 서브트리 노드 수 계산 시작
root_node = find_root(1)
dfs(root_node)

# 각 깊이(depth)별 노드 위치 기록
# tree_depth[d] = [(node 번호, 위치)] 리스트
tree_depth = [[] for _ in range(N+1)]

# 중위 순회를 하며 각 노드의 열(column) 위치 계산
def find_pos(node, depth, left_padding):
    if node != -1:
        # 중위 순회: 왼 → 루트 → 오
        # 왼쪽 먼저
        find_pos(conn[node][LEFT], depth + 1, left_padding)

        # 현재 노드 위치 = 왼쪽 서브트리 노드 수 + padding + 1
        position = left_padding + child_nums[node][LEFT] + 1
        tree_depth[depth].append((node, position))

        # 오른쪽으로 이동 (padding 증가)
        find_pos(conn[node][RIGHT], depth + 1, position)

# 중위 순회를 통해 위치 기록 시작
find_pos(root_node, 1, 0)

# 정답 초기화
ans_depth = -1
ans_width = -1

# 각 깊이마다 너비 계산 (최우측 열 - 최좌측 열 + 1)
for i in range(1, N+1):
    if tree_depth[i]:
        width = tree_depth[i][-1][1] - tree_depth[i][0][1]  # 마지막 - 처음 위치
        if ans_width < width:
            ans_width = width
            ans_depth = i

# 출력: 가장 너비가 큰 레벨과 그 너비 (1을 더해야 실제 너비가 됨)
print(ans_depth, ans_width + 1)

나는 DFS를 두번 사용하여 크기계산, 위치계산을 각각하였다.

꽤나 구현이 복잡했다.

 

GPT가 제시한 중위순위 기반 풀이도 첨부한다. 중위순위를 기반으로 크기를 1씩증가시켜 위치를 파악한다.

import sys
sys.setrecursionlimit(10**6)  # 재귀 깊이 제한 증가
input = sys.stdin.readline

N = int(input())  # 노드 개수 입력

# 트리 구조 저장용 딕셔너리
left = dict()     # left[node] = 왼쪽 자식 노드 번호
right = dict()    # right[node] = 오른쪽 자식 노드 번호
parent = [0] * (N + 1)  # 각 노드의 부모 노드 저장 (루트 찾기용)

# 트리 정보 입력
for _ in range(N):
    node, l, r = map(int, input().split())
    left[node] = l
    right[node] = r
    if l != -1:
        parent[l] = node
    if r != -1:
        parent[r] = node

# 루트 노드 찾기 (부모가 없는 노드가 루트)
root = 1
while parent[root] != 0:
    root = parent[root]

# 전역 변수: 중위 순회 순서대로 x좌표를 부여
pos = 1

# 각 깊이(depth)별 가장 왼쪽/오른쪽 x좌표 저장
level_min = dict()  # level_min[depth] = 가장 왼쪽 노드의 x좌표
level_max = dict()  # level_max[depth] = 가장 오른쪽 노드의 x좌표

# 중위 순회: 왼쪽 → 현재 노드 → 오른쪽
def inorder(node, depth):
    global pos
    if node == -1:
        return
    
    # 왼쪽 서브트리 방문
    inorder(left[node], depth + 1)

    # 현재 노드 방문 → x좌표를 pos로 지정
    if depth not in level_min:
        level_min[depth] = pos
        level_max[depth] = pos
    else:
        level_min[depth] = min(level_min[depth], pos)
        level_max[depth] = max(level_max[depth], pos)
    
    pos += 1  # 다음 노드의 x좌표

    # 오른쪽 서브트리 방문
    inorder(right[node], depth + 1)

# 루트 노드부터 중위 순회 시작 (depth 1)
inorder(root, 1)

# 레벨별 너비 계산
max_width = 0
max_depth = 0

for d in level_min:
    width = level_max[d] - level_min[d] + 1
    if width > max_width:
        max_width = width
        max_depth = d

# 결과 출력: 가장 넓은 너비를 가진 레벨(depth)과 너비
print(max_depth, max_width)