ABC074 の D 問題を python で解いてみようとしたら解けなかった

ワーシャル＝フロイド法を用いた問題演習として, AtCoder Beginner Contest 074 の D 問題の解説をしてみます.

問題はこちら
abc074.contest.atcoder.jp
これは制限時間内に解けませんでした.
問題としては重み付き無向グラフの最短距離の問題です. 特徴は既に全2点間の最短距離が与えられており,
1. そのようなグラフは存在するか？
2. 存在したとしたら, もっとも素なグラフはどんなものか？
を求める必要がありそうです. 2. はどのような意味かと言うと, 別の頂点を経由して最短距離が実現できる場合は, その辺がないグラフを考えよということです. 例えば, 入力例を考えると, 頂点から頂点への辺が不要になります.

まず, 1. そのようなグラフは存在するか？を考えます.
存在しない場合どのような状況かを考えると, 入力例のように, 他の頂点を経由した道の方が最短距離を実現できてしまう場合です. つまり, 入力に対して全2点最短距離を求め, それがと異なっていたらグラフは存在しないことになります. したがって全2点最短距離を求める際にワーシャル＝フロイド法を用い, 出力が異なっている場合はを出力すればよいです.

2. グラフが存在した場合, もっとも素なグラフにおける道の長さの和を求めます. 全ての辺について, それが必要か否かを調べます. を任意に取ります. そのとき, ある が存在して
\begin{equation}
A\[i\]\[j\]=A\[i\]\[k\]+A\[k\]\[j\]
\end{equation}
となった場合, 頂点から頂点への最短距離は頂点を経由することで達成できるためこの辺は不要になります.
一方, 任意の に対して
\begin{equation}
A\[i\]\[j\] < A\[i\]\[k\]+A\[k\]\[j\]
\end{equation}
となる場合は直接辺で結ばれないと最短距離は達成されないため, この辺は必要です.

以上により,

def Warshall_Floyd(edges,N):
  import copy
  d=copy.deepcopy(edges)
  for k in range(N):
    for i in range(N):
      for j in range(N):
        d[i][j]=min(d[i][j],d[i][k]+d[k][j])
  return d

import sys
N=int(input())
A=[list(map(int,input().split())) for i in range(N)]
B=Warshall_Floyd(A,N)
if A!=B:
  print(-1)
  sys.exit()

ans=0
for i in range(N):
  for j in range(i):
    need=1
    for k in range(N):
      if k!=i and k!=j:
        if A[i][j]==A[i][k]+A[k][j]: need=0
    ans+=need*A[i][j]
print(ans)

でACできそうです. ここでワーシャル＝フロイド法の関数内ではdeepcopyを使うことに注意しましょう. なぜならpythonでは関数の引数も参照渡しで行っているため, これをしないと緩和を行うときに元のデータ(今の場合では)が更新されてしまいます.
ところが, 残念ながらTLEします:
Submission #1605977 - AtCoder Beginner Contest 074

ワーシャル＝フロイド法などで3重ループが含まれているので, pythonでは計算が間に合いません... 一方で C++ で安直に書き直すと余裕で間に合います:
Submission #1605977 - AtCoder Beginner Contest 074

というわけでプログラミング言語に起因する要因によってACできない事態です. pythonでのforループに時間がかかってしまう話題については調べられたら記事にしようかと思います. AtCoder では PyPy という一度 C 言語に翻訳するというコンパイラがあるので, 実はこれを使うとACできるようです:
Submission #1605936 - AtCoder Beginner Contest 074