結果

問題 No.1094 木登り / Climbing tree
ユーザー ThetaTheta
提出日時 2024-03-29 11:27:12
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 3,079 bytes
コンパイル時間 355 ms
コンパイル使用メモリ 81,572 KB
実行使用メモリ 241,164 KB
最終ジャッジ日時 2024-03-29 11:28:07
合計ジャッジ時間 9,783 ms
ジャッジサーバーID
（参考情報） 		judge12 / judge11
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 60 ms
67,856 KB
testcase_01 TLE -
testcase_02 AC 609 ms
241,164 KB
testcase_03 AC 633 ms
85,508 KB
testcase_04 AC 794 ms
134,368 KB
testcase_05 AC 1,206 ms
191,496 KB
testcase_06 AC 1,188 ms
122,308 KB
testcase_07 TLE -
testcase_08 TLE -
testcase_09 TLE -
testcase_10 TLE -
testcase_11 TLE -
testcase_12 TLE -
testcase_13 TLE -
testcase_14 TLE -
testcase_15 AC 1,118 ms
109,944 KB
testcase_16 AC 1,634 ms
208,524 KB
testcase_17 AC 1,369 ms
149,764 KB
testcase_18 AC 1,249 ms
130,708 KB
testcase_19 AC 1,506 ms
184,484 KB
testcase_20 TLE -
testcase_21 AC 1,419 ms
156,332 KB
testcase_22 TLE -
testcase_23 TLE -
testcase_24 TLE -
testcase_25 TLE -
testcase_26 TLE -
権限があれば一括ダウンロードができます

ソースコード

diff # 

from math import inf, isinf
from heapq import heappop, heappush
from functools import cache
from collections import defaultdict, deque
from collections import deque
from itertools import count
from math import inf
import sys
from typing import List


def printe(*args, end="\n", **kwargs):
    print(*args, end=end, file=sys.stderr, **kwargs)


def doubling_on_tree(parent_list: List[int]) -> List[List[int]]:
    root_idx = -1
    for idx, p in enumerate(parent_list):
        if idx == p:
            root_idx = idx
            break
    else:
        raise ValueError
    double_list = [[] for _ in range(len(parent_list))]
    for idx, p in enumerate(parent_list):
        double_list[idx].append(p)
    unroot = set(range(len(parent_list)))
    for k in count(1):
        for idx in range(len(parent_list)):
            double_list[idx].append(
                double_list[double_list[idx][k - 1]][k - 1])
            if double_list[idx][-1] == root_idx:
                unroot.discard(idx)
        if not unroot:
            break
    return double_list


def calc_lca(doubling: List[List[int]],
             depth: List[int],
             idx1: int,
             idx2: int) -> int:

    if depth[idx1] < depth[idx2]:
        idx1, idx2 = idx2, idx1

    depth_diff = depth[idx1] - depth[idx2]
    shift_ctr = 0
    while depth_diff > 0:
        if depth_diff & 1:
            idx1 = doubling[idx1][shift_ctr]
        depth_diff >>= 1
        shift_ctr += 1

    if idx1 == idx2:
        return idx1

    k = len(doubling[idx1]) - 1
    while k >= 0:
        if doubling[idx1][k] != doubling[idx2][k]:
            idx1 = doubling[idx1][k]
            idx2 = doubling[idx2][k]
        k -= 1

    assert doubling[idx1][0] == doubling[idx2][0]
    return doubling[idx1][0]


def dijkstra(graph: list[dict[int, int]], start: int,
             dist_max=int(1e9)) -> list[int]:
    distance = [dist_max] * len(graph)
    queue = [(0, start)]
    distance[start] = 0
    while queue:
        c_d, c_n = heappop(queue)
        if distance[c_n] < c_d:
            continue
        for n_n, n_d in graph[c_n].items():
            if distance[n_n] > c_d + n_d:
                distance[n_n] = c_d + n_d
                heappush(queue, (distance[n_n], n_n))
    return distance


def main():
    N = int(input())
    graph = [{} for _ in range(N)]

    for _ in range(N - 1):
        a, b, c = map(int, input().split())
        graph[a - 1][b - 1] = c
        graph[b - 1][a - 1] = c
    parents = [0] * N
    depth = [inf] * N
    depth[0] = 0
    queue = deque((0, ))
    while queue:
        c = queue.popleft()
        for n in graph[c]:
            if depth[n] > depth[c]:
                depth[n] = depth[c] + 1
                parents[n] = c
                queue.append(n)

    dists = dijkstra(graph, 0)
    doubles = doubling_on_tree(parents)

    for _ in range(int(input())):
        s, t = map(lambda n: int(n) - 1, input().split())
        lca = calc_lca(doubles, depth, s, t)
        print(dists[s] + dists[t] - 2 * dists[lca])


if __name__ == "__main__":
    main()
0