백준 3176 - 도로 네트워크 (Python)

바로 전 문제와 거의 흡사합니다.
dp에 [노드, 최소거리, 최대거리] 이렇게 저장하는게 핵심입니다.
pypy3로 제출해야 합니다.

전체 코드

import sys
import math

sys.setrecursionlimit(100000)
input = sys.stdin.readline

n = int(input())
LOG = int(math.log2(n)) + 1
dp = [[[0, 0, 0] for __ in range(LOG)] for _ in range(n + 1)]
graph = [[] for _ in range(n + 1)]
depth = [0] * (n + 1)
depth[1] = 1

for _ in range(n - 1):
    a, b, w = map(int, input().split())
    graph[a].append((b, w))
    graph[b].append((a, w))

def dfs(current, num):
    for next, cost in graph[current]:
        if depth[next] == 0:
            depth[next] = num + 1
            dp[next][0] = [current, cost, cost]
            dfs(next, num + 1)

def set_parent():
    dfs(1, 1)
    for i in range(1, LOG):
        for j in range(1, n + 1):
            dp[j][i][0] = dp[dp[j][i - 1][0]][i - 1][0]
            dp[j][i][1] = min(dp[j][i - 1][1], dp[dp[j][i - 1][0]][i - 1][1])
            dp[j][i][2] = max(dp[j][i - 1][2], dp[dp[j][i - 1][0]][i - 1][2])
            

def lca(a, b):
    min_len = float('inf')
    max_len = 0
    
    if depth[a] < depth[b]:
        a, b = b, a

    for i in range(LOG, -1, -1):
        if depth[a] - depth[b] >= 2 ** i:
            min_len = min(min_len, dp[a][i][1])
            max_len = max(max_len, dp[a][i][2])
            a = dp[a][i][0]

    if a == b:
        print(min_len, max_len)
        return

    for i in range(LOG - 1, -1, -1):
        if dp[a][i][0] != dp[b][i][0]:
            min_len = min(min_len, dp[a][i][1], dp[b][i][1])
            max_len = max(max_len, dp[a][i][2], dp[b][i][2])
            a = dp[a][i][0]
            b = dp[b][i][0]

    min_len = min(min_len, dp[a][0][1], dp[b][0][1])
    max_len = max(max_len, dp[a][0][2], dp[b][0][2])
    
    print(min_len, max_len)

set_parent()

for _ in range(int(input())):
    a, b = map(int, input().split())
    lca(a, b)