백준 3176번: 도로 네트워크 (Python, PyPy3)

접근

바로 이전에 풀었던 LCA 2 문제와 매우 유사하지만, 이번에는 도로의 길이 및 최대 거리, 최소 거리가 추가된 문제였다. 최소공통조상 문제와 동일하게 풀게 되면, 두 도시 사이는 최소공통조상 도시를 거쳐 연결되게 된다.

최소공통조상 알고리즘 LCA 문제: 2021.06.29 - [코딩/백준 (Python)] - 백준 11438번: LCA 2 (Python, PyPy3)

백준 11438번: LCA 2 (Python, PyPy3)

이번 문제에서는 최소공통조상을 구하는 것에 더해 두 도시 사이의 최소거리와 최대거리를 구해야 한다. 따라서 처음에 tree 를 저장할 때 두 도시 사이의 거리를 저장해주고, 거기에 더해 dp[i][j] 가 기존에는 i 의 2^j 조상 노드를 표시했다면 이제는 dp에 [i 의 2^k 조상 도시, 해당 조상 도시까지의 최소 거리, 해당 조상 도시까지의 최대 거리] 의 리스트를 저장해준다.

이 dp를 이전 문제와 동일하게 초기화 및 희소 테이블을 계산해주고, 마찬가지로 비슷한 방식으로 두 도시 사이의 거리를 구할 수 있다.

코드

import sys
from collections import deque
from math import log2

INF = float('inf')
n = int(sys.stdin.readline())
tree = [[] for _ in range(n + 1)]
depth = [0] * (n + 1)
for _ in range(n - 1):
    a, b, c = map(int, sys.stdin.readline().split())
    tree[a].append([b, c])
    tree[b].append([a, c])

parent = [[0, 0] for _ in range(n + 1)]
check = [False] * (n + 1)
q = deque([1])
check[1] = True
while q:
    now = q.popleft()
    for b, c in tree[now]:
        if not check[b]:
            q.append(b)
            depth[b] = depth[now] + 1
            check[b] = True
            parent[b][0] = now
            parent[b][1] = c

logN = int(log2(n) + 1)
dp = [[[0, 0, 0] for _ in range(logN)] for _ in range(n + 1)]
for i in range(n + 1):
    dp[i][0][0] = parent[i][0]
    dp[i][0][1] = parent[i][1]
    dp[i][0][2] = parent[i][1]

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

k = int(sys.stdin.readline())
for _ in range(k):
    a, b = map(int, sys.stdin.readline().split())
    if depth[a] < depth[b]:
        a, b = b, a
    diff = depth[a] - depth[b]
    shortest = INF
    longest = 0
    for i in range(logN):
        if diff & 1 << i:
            shortest = min(shortest, dp[a][i][1])
            longest = max(longest, dp[a][i][2])
            a = dp[a][i][0]
    if a == b:
        print(shortest, longest)
        continue
    for i in range(logN - 1, -1, -1):
        if dp[a][i][0] != dp[b][i][0]:
            shortest = min(shortest, dp[a][i][1], dp[b][i][1])
            longest = max(longest, dp[a][i][2], dp[b][i][2])
            a = dp[a][i][0]
            b = dp[b][i][0]
    shortest = min(shortest, dp[a][0][1], dp[b][0][1])
    longest = max(longest, dp[a][0][2], dp[b][0][2])
    print(shortest, longest)

더 생각해 볼 것?

...

코드나 내용 관련 조언, 부족한 점 및 질문 언제든 환영합니다!