[BAEKJOON] 1967 트리의 지름 - Python

문제 요약

• 알고리즘 분류: 트리, DFS
• 난이도: Gold4
• 문제내용:
  • 노드와 가중치가 주어진다. 노드간의 가장 긴 길이를 구해라.

• 사이트: https://www.acmicpc.net/problem/1967

1967번: 트리의 지름

문제풀이

   이번에는 문제 유형은 트리와 DFS을 요구하는 문제이다. 트리와 DFS 알고리즘에 대한 설명은 아래에 있으니 참조하면된다.

트리 : https://jih3508.tistory.com/87

[알고리즘 이론] 트리(Tree)

DFS: https://jih3508.tistory.com/94

[알고리즘 이론] 깊이 우선 탐색(DFS)

문제 접근 방법

 일단 아래 코드 처럼 작성한면 시간초과로 뜰것이다.

import sys
input = sys.stdin.readline

def dfs(node, value):
    global max_weight
    for item in tree[node]:
        if visited[item[0]]:
            max_weight = max(max_weight, value)
        else:
            visited[item[0]] = True
            dfs(item[0], value + item[1])

N = int(input())
tree = {i : [] for i in range(1, N + 1)}
for _ in range(N - 1):
    a, b, weight = map(int, input().split())
    tree[a].append((b, weight))
    tree[b].append((a, weight))
    
max_weight = 0
for i in range(1, N + 1):
    if len(tree[i]) == 1:
        visited = [False] * (N + 1)
        dfs(i, 0)
        
print(max_weight)

 시간 초과가 나오는 이유는 각 노드끼리의 거리를 구하면 O(N ^ 2)이상 나온다. 위 문제의 데이터 개수는 10만인데 O(N ^ 2)으로 통과가 안되서 O(N)으로 처리하는 방법을 찾아야 된다. 그래서 트리와 DFS 알고리즘을 안다고해서 풀수 있는 문제가 아니다. 그래서 어느정도의 아이디어는 필요하다.

 위 문제 시간 초과를 해결하기전에 생각 해보면 노드 끼리의 연결하는 것중 가장 긴것은 리프 노드간의 연결이다. 그래서 리프 노드간의 연결중 가장 긴것을 구하면 된다. 하지만 이 방법도 O(N  / 2  * N)나오기 때문에 시간이 초과가 된다.  시간 줄이기 위해서는 말단 노드중 가장 긴 것을 구하기만 하면 O(N)이 나올것이다. 

 루트 1에서 각 리프 노트의  거리를 구하면된다. 그 중 거리가 가장 긴 리프 노드를 선택해서 긴 리프 노드에서 출발해서 각 노드의 거리를 구한 다음 큰 값을 구하면된다. 트리에서 보면 루트는 리프노드로 부터 중심이 된다. 중심점에 가장 멀리 떨어진 노드는 다른 리프노드에서 출발하면 가장 긴 노드가 될수있다는 것이다. 그럼 가장 멀리떨어진 노드를 구하면 또 그 노드에서 가장 멀리떨어진 노드의 거리를 구하면 트리의 지름이 된다.

정리

1.  루트노드에서 가장 먼 노드를 DFS로 구한다.
2. 가장 먼 노드를 기전에서 DFS로 구한 값중 가장 큰 값을 출력한다.

Code

Python

 이번 문제는 재귀 호출이 많이 이루어져서 setrecursionlimit 이것 안사용하면 에러가 날것이다.

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**5)

def dfs(node, value):
    for item in tree[node]:
        if distance[item[0]] == -1:
            distance[item[0]] = value + item[1]
            dfs(item[0], value + item[1])

N = int(input())
tree = {i : [] for i in range(1, N + 1)}
for _ in range(N - 1):
    a, b, weight = map(int, input().split())
    tree[a].append((b, weight))
    tree[b].append((a, weight))
    
distance = [-1] * (N + 1)
distance[1] = 0
dfs(1, 0)
start = distance.index(max(distance))
distance = [-1] * (N + 1)
distance[start] = 0
dfs(start, 0)
print(max(distance))