Computer Science/Algorithms

[트리] 트리의 지름 구하기

kimdozzi 2023. 9. 22. 22:17

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

 

1967번: 트리의 지름

파일의 첫 번째 줄은 노드의 개수 n(1 ≤ n ≤ 10,000)이다. 둘째 줄부터 n-1개의 줄에 각 간선에 대한 정보가 들어온다. 간선에 대한 정보는 세 개의 정수로 이루어져 있다. 첫 번째 정수는 간선이 연

www.acmicpc.net

 

트리의 지름은 DFS 두 번으로 O(N)의 시간 복잡도로 구현이 가능하다.

  1. 트리에서 아무 노드나 잡고 그 노드에 대한 먼 노드를 구하고 이 노드를 n1이라고 하자.
  2. n1에 대한 가장 먼 노드를 한번 더 구한다. 이 노드를 n2라고 하자.
  3. 이제 n1과 n2의 거리가 트리의 지름이 된다.
import sys, heapq
from collections import deque
si = sys.stdin.readline
sys.setrecursionlimit(10**8)

n = int(si())
graph = [[] for _ in range(n+1)]

def dfs(x,wei) :
    for i in graph[x] :
        a, b = i
        if distance[a] == -1 :
            distance[a] = b + wei
            dfs(a, b + wei)

for _ in range(n-1) :
    a,b,c=map(int,si().split())
    graph[a].append([b,c])
    graph[b].append([a,c])

distance = [-1] * (n+1)
distance[1] = 0
dfs(1,0)

n1 = distance.index(max(distance))
distance = [-1] * (n+1)
distance[n1] = 0
dfs(n1, 0)

print(max(distance))