結果

問題 No.1442 I-wate Shortest Path Problem
ユーザー lam6er
提出日時 2025-04-16 15:59:10
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 4,470 bytes
コンパイル時間 526 ms
コンパイル使用メモリ 82,504 KB
実行使用メモリ 166,116 KB
最終ジャッジ日時 2025-04-16 16:01:26
合計ジャッジ時間 16,296 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 7 WA * 8 TLE * 1 -- * 9
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
import heapq
from math import log2, ceil
from collections import deque

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

    N, K = int(input[ptr]), int(input[ptr+1])
    ptr += 2

    # Build the tree adjacency list
    tree = [[] for _ in range(N+1)]
    for _ in range(N-1):
        a = int(input[ptr])
        b = int(input[ptr+1])
        c = int(input[ptr+2])
        ptr += 3
        tree[a].append((b, c))
        tree[b].append((a, c))

    # Preprocess LCA and sum_from_root
    LOG = ceil(log2(N)) if N > 1 else 1
    up = [[-1]*(N+1) for _ in range(LOG)]
    sum_from_root = [0]*(N+1)
    depth = [0]*(N+1)

    # BFS to setup parent and sum_from_root
    root = 1
    up[0][root] = -1
    sum_from_root[root] = 0
    depth[root] = 0
    q = deque([root])
    visited = [False]*(N+1)
    visited[root] = True
    while q:
        u = q.popleft()
        for v, c in tree[u]:
            if not visited[v] and v != up[0][u]:
                up[0][v] = u
                sum_from_root[v] = sum_from_root[u] + c
                depth[v] = depth[u] + 1
                visited[v] = True
                q.append(v)

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

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

    def get_distance(u, v):
        ancestor = lca(u, v)
        return sum_from_root[u] + sum_from_root[v] - 2 * sum_from_root[ancestor]

    # Process airlines
    airlines = []
    min_dist = []
    for _ in range(K):
        M_i = int(input[ptr])
        P_i = int(input[ptr+1])
        ptr += 2
        X_i = list(map(int, input[ptr:ptr+M_i]))
        ptr += M_i
        airlines.append((M_i, P_i, X_i))

    # Compute min_dist for each airline using multi-source Dijkstra
    for i in range(K):
        M_i, P_i, X_i = airlines[i]
        dist = [float('inf')] * (N + 1)
        heap = []
        for s in X_i:
            if dist[s] > 0:
                dist[s] = 0
                heapq.heappush(heap, (0, s))
        while heap:
            d, u = heapq.heappop(heap)
            if d > dist[u]:
                continue
            for v, cost in tree[u]:
                if dist[v] > d + cost:
                    dist[v] = d + cost
                    heapq.heappush(heap, (dist[v], v))
        min_dist.append(dist)

    # Precompute sum_p for all subsets
    sum_p = [0] * (1 << K)
    for mask in range(1, 1 << K):
        s = 0
        for i in range(K):
            if mask & (1 << i):
                s += airlines[i][1]
        sum_p[mask] = s

    # Process queries
    Q = int(input[ptr])
    ptr += 1
    for _ in range(Q):
        U = int(input[ptr])
        V = int(input[ptr+1])
        ptr += 2
        road_cost = get_distance(U, V)
        if K == 0:
            print(road_cost)
            continue

        d_u = [min_dist[i][U] for i in range(K)]
        d_v = [min_dist[i][V] for i in range(K)]

        # Compute min_mask_u
        min_mask_u = [float('inf')] * (1 << K)
        for mask in range(1, 1 << K):
            lsb = mask & -mask
            i = (lsb).bit_length() - 1
            prev_mask = mask ^ (1 << i)
            min_mask_u[mask] = min(min_mask_u[prev_mask], d_u[i]) if prev_mask != 0 else d_u[i]

        # Compute min_mask_v
        min_mask_v = [float('inf')] * (1 << K)
        for mask in range(1, 1 << K):
            lsb = mask & -mask
            i = (lsb).bit_length() - 1
            prev_mask = mask ^ (1 << i)
            min_mask_v[mask] = min(min_mask_v[prev_mask], d_v[i]) if prev_mask != 0 else d_v[i]

        # Find minimal candidate cost
        min_candidate = float('inf')
        for mask in range(1, 1 << K):
            candidate = sum_p[mask] + min_mask_u[mask] + min_mask_v[mask]
            if candidate < min_candidate:
                min_candidate = candidate

        answer = min(road_cost, min_candidate)
        print(answer)

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