[백준 1967] 트리의 지름 - Python (트리, DFS)

 

1967번: 트리의 지름

 

 

<내 코드>

 

import sys
sys.setrecursionlimit(100000)


def dfs(start, tree, weight, ck):
    # ck로 루트1번에서 시작한 건지, 최장길이 노드를 루트로 한건지 판단
    if ck == 1:
        weight[1] = 0  # 루트 1번 노드 가중치 0으로 설정

    for node, w in tree[start]:
        if(weight[node] == 0):
            weight[node] = weight[start] + w  # 현재 가중치 + 다음 노드와 가중치
            dfs(node, tree, weight, ck)


n = int(sys.stdin.readline())
tree = [[] for _ in range(n+1)]
# 입력 값 tree 생성
for _ in range(n-1):
    p_node, c_node, w = map(int, sys.stdin.readline().split())
    tree[p_node].append((c_node, w))
    tree[c_node].append((p_node, w))

weight1 = [0 for _ in range(n+1)]  # 루트1로부터 길이를 저장
dfs(1, tree, weight1, 1)  # 루트 1 시작점으로 해 가장 먼 노드를 찾음


# 찾은 가장 먼 노드를 시작 노드로 탐색
# 최장경로 노드에서 다시 DFS를 통해 트리지름구하기
start_node = weight1.index(max(weight1))
weight2 = [0 for _ in range(n+1)]

dfs(start_node, tree, weight2, 2)

print(max(weight2))

같은 코드로 python3, pypy3로 제출했을 때 시간 차이가 5000ms 정도 차이가 났다...

확실히 DFS는 시간이 많이 걸리는 것 같다. 파이썬은 테스트에서 되도록이면 BFS를 이용하는 게 좋지 않을까 생각한다.

 

알고리즘 이해

: 일단 트리의 지름을 구하려면 가장 끝 노드끼리 거리가 될거다. 거기서도 가장 큰 가중치(거리)를 가지는 것이 최종 지름이 된다.

 

1. 루트 1번을 시작으로 해서 가장 긴(먼 거리)에 있는 즉, 가중치가 큰  끝 노드를 DFS 탐색으로 찾는다.

2. 찾은 끝 노드를 다시 루트노드로 설정해 DFS 탐색을 한다.

3. 탐색하며 더한 가중치 값들 중 최댓값을 출력하면, 트리의 지름이 된다.