[백준/파이썬]1967 트리의 지름

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

1967번: 트리의 지름

문제

 

백준 1967번 트리의 지름

 

문제는 두 노드를 선택했을 때, 두 노드 사이의 지름의 길이가 최대가 되는 길이를 출력하는 문제다.

 

처음 문제를 풀기 위해서 접근한 방법은 두 노드 쌍을 모두 구하고, 여기에 대해서 DFS탐색을 통해 가장 긴 지름을 탐색해 봤다.

 

구현하면서도 너무 오래 걸릴 거라 생각했는데, 역시 시간초과가 났다.

 

오답

import sys
sys.setrecursionlimit(2500)
input = sys.stdin.readline
def dfs(node, goal, value) :
    global ans
    if node == goal :
        if ans < value :
            ans = value
        return
    visit[node] = True
    for j in range(len(tree[node])) :
        if visit[tree[node][j][0]] == False :
            dfs(tree[node][j][0], goal, value+tree[node][j][1])
        

n = int(input())

tree = [[] for _ in range(n+1)]

for i in range(n-1) :
    s, e, v = map(int, input().split())

    tree[s].append([e, v])
    tree[e].append([s, v])

ans = 0

for i in range(1, n+1) :
    for k in range(1, n+1) :
        if i == k : continue

        visit = [False for _ in range(n+1)]

        dfs(i, k, 0)

 

그래서 두 번째로 진행한 방법은 모든 쌍을 구하는 것이 아닌, 각 노드를 기준으로 가장 긴 지름을 찾는 방법을 선택했다.

 

그런데 이것도 모든 쌍을 구하는 연산만 빼면 시간 복잡도가 거의 비슷하다.

 

그래서 역시 이것도 시간 초과가 났다.

 

 

그래서 GPT의 약간의 도움을 받았다.

 

GPT왈 루트 노드를 기준으로 가장 길이가 긴(멀리) 있는 노드를 찾고 그 노드를 기준으로 가장 긴 지름을 찾으라는 것이었다.

 

이렇게 하면 각 노드에 대해서 DFS연산을 하지 않아도 되기 때문에 시간 복잡도가 크게 작아지게 된다. 다음은 이 방식대로 구현한 코드다.

 

정답

import sys
sys.setrecursionlimit(1000000)
input = sys.stdin.readline
def dfs(node, value) :
    global ans, start_node

    if ans < value :
        ans = value
        start_node = node
    visit[node] = True
    for j in range(len(tree[node])) :
        if visit[tree[node][j][0]] == False :
            dfs(tree[node][j][0], value+tree[node][j][1])
        

n = int(input())
if n==1 :
    print("0")
    exit()
tree = [[] for _ in range(n+1)]

for i in range(n-1) :
    s, e, v = map(int, input().split())

    tree[s].append([e, v])
    tree[e].append([s, v])

ans = 0

visit = [False for _ in range(n+1)]
start_node = None
dfs(1, 0)

ans = 0
visit = [False for _ in range(n+1)]
dfs(start_node, 0)

print(ans)

 

코드에서 알 수 있듯, 첫 번째 DFS탐색에서는 루트 노드로부터 가장 멀리 있는 노드를 탐색한다.

 

멀리 있는 노드를 찾았으면, 그 노드를 기준으로 지름이 가장 긴 길이를 찾는다. 이렇게 하면 두 번의 DFS탐색으로 문제 해결이 가능하다.

---

 

 

 GPT 녀석 돈 주니 더 똑똑해졌다. 문제 풀다가 안 풀리면 애용해야겠다.