Python - [백준]1967-트리의 지름

출처

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

문제

노드를 잇는 간선의 가중치를 주고, 그 트리의 가장 긴 길이인 지름을 찾는 문제이다.

접근 방법

Dfs알고리즘을 사용하여 어떤 리프 노드로 부터 다른 리프노드에 도달할 때 까지의 최대값을 출력하도록 하는 문제이다.
메모리 초과가 나기때문에 N*N형식의 배열을 사용하여 가중치를 기록 할 순 없다. 그래서 3차원 배열을 활용하여 map[1] = [ [이어진 노드, 가중치] 들 ] - 1번으로 부터 이어진 노드들의 형식을 사용하였다.

풀이

처음엔 모든 리프노드에 대해 dfs를 하여 모든 경우를 탐색하려고 하였지만 메모리 초과에 걸렸다. 고민을 해보니 다음과 같은 경우를 생각해냈다.

• 루트노드로 부터 가장 먼 리프노드를 찾고 (DFS)
• 그 리프노드로 부터 가장 먼 리프노드를 찾아낸다(DFS)

총 2번의 DFS를 적용시켜 해결하여 메모리 초과를 해결하였다.
사실 이 개념을 떠올리기까지 오래 걸렸지만 떠올리고 나면 구현하기 어려운 문제는 아니라고 생각한다..

코드

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**9)
n = int(input())

load = [[] for i in range(n+1)]
answer = 0
for i in range(n-1):
    a, b, c = map(int, input().split())
    load[a].append((b, c))
    load[b].append((a, c))

leaf = [-1 for _ in range(n+1)]
result = [-1 for _ in range(n+1)]
def dfs(x, array):
    for i in load[x]:
        if array[i[0]] == -1:
            array[i[0]] = array[x] + i[1]
            dfs(i[0],array)
leaf[1] = 0
dfs(1, leaf)
idx = 0
for i in range(n+1):
    if leaf[i] == max(leaf):
        idx = i
result[idx] = 0
dfs(idx, result)
print(max(result))