結果

問題 No.1212 Second Path
ユーザー lam6er
提出日時 2025-04-09 21:04:25
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 2,925 bytes
コンパイル時間 246 ms
コンパイル使用メモリ 82,300 KB
実行使用メモリ 168,416 KB
最終ジャッジ日時 2025-04-09 21:06:11
合計ジャッジ時間 12,515 ms
ジャッジサーバーID
(参考情報)
judge5 / judge2
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample -- * 3
other TLE * 1 -- * 44
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys
sys.setrecursionlimit(1 << 25)

from math import log2, ceil
from collections import deque

def main():
    input = sys.stdin.read().split()
    ptr = 0

    N = int(input[ptr]); ptr +=1

    edges = []
    adj = [[] for _ in range(N+1)]  # nodes are 1-based

    for _ in range(N-1):
        u = int(input[ptr]); ptr +=1
        v = int(input[ptr]); ptr +=1
        w = int(input[ptr]); ptr +=1
        adj[u].append((v, w))
        adj[v].append((u, w))
        edges.append((u, v, w))

    # LCA with binary lifting.
    LOG = ceil(log2(N)) if N > 1 else 1
    parent = [[-1]*(N+1) for _ in range(LOG)]
    depth = [0]*(N+1)
    dist = [0]*(N+1)  # distance from root (1) to node.

    # BFS for parent, depth, dist.
    q = deque()
    root = 1
    q.append(root)
    parent[0][root] = -1
    while q:
        u = q.popleft()
        for v, w in adj[u]:
            if parent[0][v] == -1 and v != root:
                parent[0][v] = u
                depth[v] = depth[u] + 1
                dist[v] = dist[u] + w
                q.append(v)

    # Build binary lifting table.
    for k in range(1, LOG):
        for v in range(1, N+1):
            if parent[k-1][v] != -1:
                parent[k][v] = parent[k-1][parent[k-1][v]]
            else:
                parent[k][v] = -1

    # LCA function.
    def lca(u, v):
        if depth[u] < depth[v]:
            u, v = v, u
        # Bring u up to depth v.
        for k in reversed(range(LOG)):
            if depth[u] - (1 << k) >= depth[v]:
                u = parent[k][u]
        if u == v:
            return u
        # Now find the LCA.
        for k in reversed(range(LOG)):
            if parent[k][u] != -1 and parent[k][u] != parent[k][v]:
                u = parent[k][u]
                v = parent[k][v]
        return parent[0][u]

    # Distance function.
    def get_dist(u, v):
        return dist[u] + dist[v] - 2 * dist[lca(u, v)]

    # Function to check if u is on the path between x and y.
    def is_on_path(x, y, u):
        l = lca(x, y)
        if lca(x, u) == u and lca(u, y) == l:
            return True
        if lca(y, u) == u and lca(u, x) == l:
            return True
        return False

    # Sort edges by weight.
    edges_sorted = sorted(edges, key=lambda x: x[2])

    Q = int(input[ptr]); ptr +=1

    for _ in range(Q):
        x = int(input[ptr]); ptr +=1
        y = int(input[ptr]); ptr +=1

        if x == y:
            print(-1)
            continue

        S = get_dist(x, y)
        found = False

        for u, v, w in edges_sorted:
            # Check if exactly one of u or v is on the path.
            on_u = is_on_path(x, y, u)
            on_v = is_on_path(x, y, v)
            if (on_u and not on_v) or (on_v and not on_u):
                print(S + 2 * w)
                found = True
                break
        if not found:
            print(-1)

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