結果

問題 No.2604 Initial Motion
コンテスト
ユーザー detteiuu
提出日時 2024-12-03 23:27:07
言語 PyPy3
(7.3.15)
結果
MLE  
実行時間 -
コード長 2,760 bytes
コンパイル時間 546 ms
コンパイル使用メモリ 82,300 KB
実行使用メモリ 1,575,084 KB
最終ジャッジ日時 2024-12-03 23:29:19
合計ジャッジ時間 127,862 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1 MLE * 2
other AC * 8 TLE * 8 MLE * 23
権限があれば一括ダウンロードができます

ソースコード

diff # 

from heapq import heappush, heappop
class Edge:
    def __init__(self, to, cap, cost, rev):
        self.to = to
        self.cap = cap
        self.cost = cost
        self.rev = rev

class MinCostFlow:
    def __init__(self, N):
        self.N = N
        self.G = [[] for i in range(N)]

    def add_edge(self, fr, to, cap, cost):
        forward = [to, cap, cost, None]
        backward = forward[3] = [fr, 0, -cost, forward]
        self.G[fr].append(Edge(to, cap, cost, len(self.G[to])))
        self.G[to].append(Edge(fr, 0, -cost, len(self.G[fr])-1))

    def flow(self, s, t, f):
        ans = 0
        H = [0]*self.N
        prv_v = [0]*self.N
        prv_e = [None]*self.N
        dist = [INF]*self.N
        while f:
            dist = [INF]*self.N
            dist[s] = 0
            que = [(0, s)]
            while que:
                d, now = heappop(que)
                if dist[now] < d:
                    continue
                r0 = dist[now] + H[now]
                for e in self.G[now]:
                    if e.cap > 0 and r0 + e.cost - H[e.to] < dist[e.to]:
                        dist[e.to] = r0 + e.cost - H[e.to]
                        prv_v[e.to] = now; prv_e[e.to] = e
                        heappush(que, (dist[e.to], e.to))
            if dist[t] == INF:
                return -1
            for i in range(self.N):
                H[i] += dist[i]
            d = f
            v = t
            while v != s:
                d = min(d, prv_e[v].cap)
                v = prv_v[v]
            f -= d
            ans += d * H[t]
            v = t
            while v != s:
                e = prv_e[v]
                e.cap -= d
                self.G[e.to][e.rev].cap += d
                v = prv_v[v]
        return ans

INF = 10**18
def dijkstra(start):
    dist = [INF]*N
    dist[start] = 0
    visited = [False]*N
    que = [(0, start)]
    while que:
        d, now = heappop(que)
        if visited[now]:
            continue
        visited[now] = True
        for next, weight in G[now]:
            if dist[now]+weight < dist[next]:
                dist[next] = dist[now]+weight
                heappush(que, (dist[next], next))
    return dist

K, N, M = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
G = [[] for _ in range(N)]
for _ in range(M):
    U, V, D = map(int, input().split())
    G[U-1].append((V-1, D))
    G[V-1].append((U-1, D))

dist = []
for i in range(N):
    dist.append(dijkstra(i))

MCF = MinCostFlow(K+N+2)
for i in range(K):
    MCF.add_edge(K+N, i, 1, 0)
for i in range(N):
    MCF.add_edge(K+i, K+N+1, B[i], 0)
for i in range(K):
    for j in range(N):
        MCF.add_edge(i, K+j, 1, dist[A[i]-1][j])

print(MCF.flow(K+N, K+N+1, K))
0