BOJ 1967번: 트리의 지름

www.acmicpc.net/problem/1967

1967번: 트리의 지름

아이디어:

트리의 지름에 대한 증명이 있다.

 1. 트리에서 임의의 정점 x를 잡는다.

 2. 정점 x에서 가장 먼 정점 y를 찾는다.

 3. 정점 y에서 가장 먼 정점 z를 찾는다.

    -> y-z의 경로가 트리의 지름이다.

 

이 알고리즘? 정의? 를 이용해서 풀어보자.

DFS와 재귀를 이용하여 구현했는데, 답은 제대로 나오는데 시간초과가 난다. pypy로 제출해도..ㅠㅠ

import sys
sys.setrecursionlimit(10**9)

n = int(sys.stdin.readline())
graph = [[] for _ in range(n+1)]

for _ in range(n-1):
    parent, child, val = map(int, sys.stdin.readline().split())
    graph[parent].append([child, val])
    graph[child].append([parent, val])
    
visit = [False]*(n+1)
temp = 0
ans = 0
y = 0
def DFS(x):
    global ans, y, temp
    if (ans<temp):
        ans = temp
        y = x
    
    for i in graph[x]:
        if (visit[i[0]]==False):
            visit[i[0]] = True
            temp+=i[1]
            DFS(i[0])
            visit[i[0]] = False
            temp-=i[1]
DFS(1)
visit = [False]*(n+1)
temp = 0
DFS(y)
print(ans)

 

구글링을 해보니까 나처럼 temp로 계속 더했다 뺐다 계산하지 않고 그냥 루트부터 해당하는 노드까지의 거리를 배열을 따로 만들어서 저장하고 DFS탐색을 마친 후에 배열의 최댓값을 구하는 방식으로 풀더라. 그리고 이렇게 푸니까 시간초과가 해결됐다..ㅠㅠ

아무래도 비교연산때문에 시간초과가 난 게 아닐까 생각한다.

 

아래가 통과된 코드이다.

import sys
sys.setrecursionlimit(10**9)

n = int(sys.stdin.readline())
graph = [[] for _ in range(n+1)]

for _ in range(n-1):
    parent, child, val = map(int, sys.stdin.readline().split())
    graph[parent].append([child, val])
    graph[child].append([parent, val])
    
visit = [False]*(n+1)

def DFS(x, result):
    visit[x] = True

    for i in graph[x]:
        
        if (visit[i[0]]==False):
            result[i[0]] = result[x]+i[1]
            DFS(i[0], result)

result1 = [0]*(n+1) #1로부터 거리
DFS(1, result1)

visit = [False]*(n+1)
index = result1.index(max(result1))
result2 = [0]*(n+1)
DFS(index, result2)
print(max(result2))