ワーシャルフロイド法でABC 208 D - Shortest Path Queries 2を解く

はんなりプロコン AtCoder部 #14 - connpass で取り組んだ問題。理解するのに時間はかかったけど、わかってしまえば簡単だった。

解説は公式よりもユーザー解説の方がわかりやすかった。

blog.hamayanhamayan.com

わかりにくかった部分

k = 0 とした場合ってつまり何？

今回はk=0から始めてみよう。（k=1からじゃないの？という突っ込みはちょっと置いておいて…）これで一旦f(s,t,0)を考えてみると、使えるのはs->tへの直接の遷移だけとなる。

k=0から始めるとあるが、その状態はどう定義すれば良いか、というところがわかりにくい部分だった。
k=0とはつまり直接遷移した時の移動コストになるので、問題でいうところのM本の道として与えられるA B Cの値になる。なので、標準入力から読み込んだ時に初期値として代入できる。

(N, M) = gets.chomp.split.map(&:to_i)
 
# 初期化 - 到達出来ない時はコスト。無限
inf = 10 ** 12
dist = Array.new(N) { Array.new(N) { inf } }

# 同じ都市への移動はコスト0
(0...N).each do |i|
  dist[i][i] = 0
end

M.times {
  # k = 0 の時は直接遷移するコストとなる。
  (s,t,cost) = gets.split.map(&:to_i)
  dist[s-1][t-1] = cost
}

この状態からk=1~N までdist 配列を操作して最短経路をもとめていけば結果が

sum=0
(1...N).each do |k|
  (0...N).each do |s|
    (0...N).each do |t|
      dist[s][t] = [dist[s][t], dist[s][k] + dist[k][t]].min
      sum += dist[s][t] unless dist[s][t] == inf
    end
  end
end

アルゴリズムは合っていてもTLEになってしまう

Rubyは動的型付のスクリプト言語なので簡単にコードを書けますが、実行速度はコンパイラ型言語に比べると遅いです。今回も上記のように書きやすいように書いたコードではTLE（実行時間制限超過）になってしまった！！

Rangeによる繰り返しはコストが高い

ワーシャルフロイド法を想定した三重ループを実装し実行時間を比較。while を使った方がRangeを使うより1.40倍ほど速かった。

Rangeを使う繰り返し

N = 400

dist = Array.new(N) { Array.new(N, 10**12) }

(1...N).each do |k|
  (0...N).each do |s|
    (0...N).each do |t|
      dist[s][t] = [dist[s][t], dist[s][k] + dist[k][t]].min
    end
  end
end

5回の実行時間の平均: 7.264秒

Whileを使う繰り返し

N = 400

dist = Array.new(N) { Array.new(N, 10**12) }

k=1
while k < N
  s=0
  while s < N
    t=0
    while t < N
      dist[s][t] = [dist[s][t], dist[s][k] + dist[k][t]].min
      t += 1
    end
    s += 1
  end
  k += 1
end

5回の実行時間の平均: 5.187秒

配列アクセスはコストが高い

経路コストを比較する際に、添字を指定して二次元配列に毎回アクセスをしている。繰り返しの中でこれを行うとコストが高くなってしまう。

dist[s][t] = [dist[s][t], dist[s][k] + dist[k][t]].min

これを変数に持ち替えて（キャッシュ）扱うように変えると、1.95倍ほど速かった

二次元配列にアクセス

先のWhileを使う繰り返しのまま

N = 400

dist = Array.new(N) { Array.new(N, 10**12) }

k=1
while k < N
  s=0
  while s < N
    t=0
    while t < N
      dist[s][t] = [dist[s][t], dist[s][k] + dist[k][t]].min
      t += 1
    end
    s += 1
  end
  k += 1
end

可能な限り変数に持つ

N = 400

dist = Array.new(N) { Array.new(N, 10**12) }

k=1
while k < N
  s=0
  distk = dist[k]
  while s < N
    t=0
    dists  = dist[s]
    distsk = dists[k]
    while t < N
      c = distsk + distk[t]
      dists[t] = c if c < dists[t]
      t += 1
    end
    s += 1
  end
  k += 1
end