aacord’s memo

abcを中心にpythonで解いた問題のメモ、整理をしています。緑になった。

abc 074 D - Restoring Road Network (python)

サンプルを使ってワーシャルフロイドした結果が、サンプルと一致していなければ答えは -1 となる。
そうでない場合は存在する道の和を出力しないといけないが、それは i から j に道が引かれているとき == min( d[ i ][ k ] + d[ k ][ j ] ) > d[ i ][ j ] で求められる。
numpy 使うと便利

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np

from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import floyd_warshall

n = int(readline())
f = []
A = []
chk = []
for i in range(n):
  ar = list(map(int,readline().split()))
  f.append(ar)
  for j,ai in enumerate(ar):
    if i == j:
      continue
    A.append((i,j,ai))
  
INF = 10**9+1
g = [[INF] * n for _ in range(n)]
for a,b,c in A:
    g[a][b] = c
    g[b][a] = c

g = csr_matrix(g)
d = floyd_warshall(csgraph=g)

t = f == d
if t.all():
  ans = 0
  for i in range(n):
    d[i,i] = INF

  for i in range(n):
    for j in range(i+1,n):
      opt = np.min(d[i] + d[j])
      if opt > d[i,j]:
        ans += d[i,j]
        
  print(int(ans))
else:
  print(-1)