結果

問題 No.1812 Uribo Road
ユーザー lam6er
提出日時 2025-03-20 20:26:49
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 611 ms / 5,000 ms
コード長 2,715 bytes
コンパイル時間 150 ms
コンパイル使用メモリ 82,484 KB
実行使用メモリ 96,280 KB
最終ジャッジ日時 2025-03-20 20:28:17
合計ジャッジ時間 8,161 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 4
other AC * 30
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
from heapq import heappush, heappop
from collections import defaultdict

INF = float('inf')

def main():
    input = sys.stdin.read().split()
    ptr = 0
    N = int(input[ptr]); ptr +=1
    M = int(input[ptr]); ptr +=1
    K = int(input[ptr]); ptr +=1
    R = list(map(lambda x: int(x)-1, input[ptr:ptr+K]))
    ptr +=K
    
    edges = []
    adj = [[] for _ in range(N+1)]
    for _ in range(M):
        A = int(input[ptr]); ptr +=1
        B = int(input[ptr]); ptr +=1
        C = int(input[ptr]); ptr +=1
        edges.append( (A, B, C) )
        adj[A].append( (B, C) )
        adj[B].append( (A, C) )
    
    required_edges = []
    for r in R:
        A, B, C = edges[r]
        required_edges.append( (A, B, C) )
    
    targets = []
    for j in range(K):
        A_j, B_j, C_j = required_edges[j]
        targets.append(A_j)
        targets.append(B_j)
    targets.append(N)
    
    target_dists = []
    for t in targets:
        dist = dijkstra(t, adj)
        target_dists.append(dist)
    
    max_mask = 1 << K
    dp = [defaultdict(lambda: INF) for _ in range(max_mask)]
    dp[0][1] = 0
    
    for mask in range(max_mask):
        current_nodes = list(dp[mask].items())
        for u, cost in current_nodes:
            if cost == INF:
                continue
            for j in range(K):
                if (mask & (1 << j)) != 0:
                    continue
                A_j, B_j, C_j = required_edges[j]
                s_index = 2*j
                d_s = target_dists[s_index][u]
                new_cost = cost + d_s + C_j
                new_mask = mask | (1 << j)
                t = B_j
                if new_cost < dp[new_mask][t]:
                    dp[new_mask][t] = new_cost
                
                s_index = 2*j + 1
                d_s = target_dists[s_index][u]
                new_cost = cost + d_s + C_j
                new_mask = mask | (1 << j)
                t = A_j
                if new_cost < dp[new_mask][t]:
                    dp[new_mask][t] = new_cost
    
    full_mask = (1 << K) - 1
    min_total = INF
    n_index = len(targets) - 1
    dist_to_N = target_dists[n_index]
    for u, cost in dp[full_mask].items():
        total = cost + dist_to_N[u]
        if total < min_total:
            min_total = total
    print(min_total)

def dijkstra(start, adj):
    n = len(adj) - 1
    dist = [INF] * (n + 1)
    dist[start] = 0
    heap = [(0, start)]
    while heap:
        d, u = heappop(heap)
        if d > dist[u]:
            continue
        for v, w in adj[u]:
            if dist[v] > d + w:
                dist[v] = d + w
                heappush(heap, (dist[v], v))
    return dist

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