백준 1967: 트리의 지름 - dfs/bfs, 트리(Python, 파이썬)

트리의 지름을 구하는 방법.

임의의 한 정점에서 가장 먼 정점 a를 구한다.
a에서 가장 먼 정점 b를 구한다.
a에서 b까지의 거리 == 트리의 지름.

처음에 다익스트라로 풀었다가 트리의 지름을 구하는 방법 알고 dfs/bfs 로 전환.

파이썬 함수 리턴값

오늘 알게 된 사실 하나 더.
파이썬 함수 리턴값 여러 개일 땐 튜플로 변환되어 나온다는 점.

dfs

import sys
input = sys.stdin.readline

def dfs(start):
    visited = [0] * (n+1)
    stack = [(0, start)]
    visited[start] = -1

    while stack:
        dist, now = stack.pop()

        for v, next in tree[now]:
            if not visited[next]:
                visited[next] = dist + v
                stack.append((dist + v, next))

    mv = max(visited)
    return mv, visited.index(mv)


n = int(input())
tree = [[] for _ in range(n+1)]
for i in range(n-1):
    p, c, v = map(int, input().split())
    tree[p].append((v, c))
    tree[c].append((v, p))

print(dfs(dfs(1)[1])[0])

dijkstra

import sys
import heapq
input = sys.stdin.readline

def dijkstra(start):
    dist = [-1] + [float('inf')]* n
    pq = [(0, start)] 
    dist[start] = 0

    while pq:
        c, now = heapq.heappop(pq)
        if c > dist[now]:
            continue

        for v, next in tree[now]:
            cost = v + c
            if dist[next] > cost:
                dist[next] = cost
                heapq.heappush(pq, (cost, next))

    max_dist = max(dist)
    idx = dist.index(max_dist)
    return max_dist, idx

n = int(input())
tree = [[] for _ in range(n+1)]
for i in range(n-1):
    p, c, v = map(int, input().split())
    tree[p].append((v, c))
    tree[c].append((v, p))

print(dijkstra(dijkstra(1)[1])[0])

bfs

import sys
from collections import deque
input = sys.stdin.readline

def bfs(start):
    visited = [0] * (n+1)
    q = deque([(0, start)])
    visited[start] = -1

    while q:
        dist, now = q.popleft()

        for v, next in tree[now]:
            if not visited[next]:
                visited[next] = dist + v
                q.append((dist + v, next))

    mv = max(visited)
    return mv, visited.index(mv)


n = int(input())
tree = [[] for _ in range(n+1)]
for i in range(n-1):
    p, c, v = map(int, input().split())
    tree[p].append((v, c))
    tree[c].append((v, p))

print(bfs(bfs(1)[1])[0])

결과는 dfs - 다익스트라 -bfs 순으로 빠르게 나왔다네요~