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

2250번: 트리의 높이와 너비

 

코드

def inorder(node,level):
	global dist
	if n_list[node][0] != -1:
		inorder(n_list[node][0],level+1)
	distance[level].append(dist)
	dist += 1
	if n_list[node][1] != -1:
		inorder(n_list[node][1],level+1)

n = int(input())
n_list = [[0,0] for _ in range(n+1)]
distance = [[] for _ in range(n+1)]
r_list = [0 for _ in range(n+1)]
dist = 1
for _ in range(n):
	parent,l,r = map(int,input().split())
	n_list[parent][0] = l
	n_list[parent][1] = r
	if l != -1:
		r_list[l] += 1
	if r != -1:
		r_list[r] += 1
for i in range(1,n+1):
	if r_list[i] == 0:
		root = i

inorder(root,1)

max_dist = 0
l = 1
for i in range(1,n+1):
	if distance[i]:
		d = max(distance[i]) - min(distance[i]) + 1
		if max_dist < d:
			l = i
			max_dist = d
print(l,max_dist)

 

설명

중위순회를 하면서 방문하는 순서가 해당 노드의 열 번호(dist)가 된다.

같은 레벨에 있는 노드들끼리 열 번호를 distance[level]에 저장한다.

for문을 이용하여 해당 레벨에서의 너비를 구하고 현재 최댓값과 비교하여 너비의 최댓값을 계속 갱신시켜주면 된다.

여기서 주의할 점은 트리의 루트가 항상 1이 아니기 때문에 처음에 루트 번호를 따로 찾아줘야 된다는 것이다.