백준 1167번 트리의 지름

백준 1167번 트리의 지름

문제 바로가기

코드

import sys

answer = 0
N = int(sys.stdin.readline())
tree = {i : [] for i in range(1, N+1)}
check = [-1]*(N+1)

for i in range(N):
  li = list(map(int, sys.stdin.readline().split()))
  root = li[0]
  li.pop()
  while len(li) > 1:
    dist = li.pop()
    node = li.pop() 
    tree[root].append((node, dist))


def findMaxIndex(node, check):
  for destination, dist in tree[node]:
    if check[destination] == -1:
      check[destination] = check[node] + dist  
      findMaxIndex(destination, check)

def dfs(node, r, visited):
  global answer
  for destination, dist in tree[node]:
    if destination not in visited:
      visited.add(node)
      dfs(destination, r+dist, visited)
      visited.remove(node)
    else:
      answer = max(answer, r)
    

findMaxIndex(1, check)
dist = 0
startNode = 1 
for i in range(1, N+1):
  if dist < check[i]:
    dist = check[i]
    startNode = i


dfs(startNode, 0, set())

print(answer)

🔧 문제 해결

이 문제는 트리가 주어지고 각 트리를 연결하는 간선의 가중치가 있는 모습이다.

여기서 한 노드에서 다른 노드까지(거쳐가도 된다) 최대 가중치의 합을 구하는 문제이다.

처음에는 이 문제를 DFS 탐색을 통해서 해결하려고 했다. 주어진 트리의 모든 노드에서 출발하여서 최대치의 가중치 합을 각각 구하여서 이중에 최대값을 리턴하는 방식으로 말이다. 하지만.... 트리의 노드 Input이 (2 ≤ V ≤ 100,000) 크므로 해당 방법으로는 시간이 너무 오래걸렸다.

시간초과에 대한 문제를 해결하기 위해서 사용한 방법은 최대값이 나올만한 노드를 미리 찾아서 해당 노드만 탐색을 진행하는 방법이다. 트리의 노드들은 모두 연결이 되어있으며 한 노드에서 다른 노드까지의 가중치의 합을 구해서 비교해보면 최대값이 나올 수 있는 노드를 구할 수 있다.

👉 여기서 한 노드만을 통해서 구한 최대값을 그대로 사용하면 되지 않는가? 라는 의문이 떠오를 수 있지만, 만약 해당 노드의 부모노드가 존재하게 된다면 이 값은 최종적인 답이 될 수 없다.

따라서 findMaxIndex를 통해서 dfs 탐색을 할 노드를 찾고 탐색을 진행해주면 된다.

👍 도움이 되었던 반례 목록

4
1 2 5 3 9 -1
2 1 5 -1
3 1 9 4 8 -1
4 3 8 -1

답 : 22

6
1 2 3 -1
2 1 3 5 3 3 5 -1
3 2 5 4 7 -1
4 3 7 -1
5 2 3 6 5 -1
6 5 5 -1

답 : 20

4
1 2 7 3 2 -1
2 1 7 -1
3 1 2 4 3 -1
4 3 3 -1

답 : 12

5
1 2 7 3 2 5 10 -1
2 1 7 -1
3 1 2 4 3 -1
4 3 3 -1
5 1 10 -1

답 : 17