多次元配列NArrayの機能を利用すれば、ワーシャル–フロイド法の3重ループのうち内側2つを書かなくてよくなる。Rubyレベルの記述が減ることで計算速度も向上する。
（AtCoderのコンテストでスラスラ書けず、NArray不使用で提出してTLEに苦しんだので、きちんと纏めておくことにした。）
コード
# n頂点m辺の重み付き有向グラフを標準入力から読み込み
n, m = gets.split.map!(&:to_i)
edges = Array.new(m) { gets.split.map!(&:to_i) }
# 隣接行列を作成
d = Numo::Int64.new(n, n).fill(Numo::Int64::MAX >> 1)
d.diagonal[true] = 0
edges.each do |i, j, dist|
d[i, j] = dist if d[i, j] > dist
end
# ワーシャル–フロイド法
n.times do |k|
d_via_k = d[true, k..k] + d[k..k, true]
d = Numo::Int64.minimum(d, d_via_k)
end
# 結果（全ペアの最短経路）を出力
p d
入出力のサンプル
Wikipedia英語版の例を拝借する。（頂点の添字は0始まりに変更）
input
4 5
0 2 -2
1 0 4
1 2 3
2 3 2
3 1 -1
output
Numo::Int64#shape=[4,4]
[[0, -1, -2, 0],
[4, 0, 2, 4],
[5, 1, 0, 2],
[3, -1, 1, 0]]
変形のメモ
ワーシャル–フロイド法の3重ループをふつうに書けばこうなる。
n.times do |k|
n.times do |i|
n.times do |j|
d_ikj = d[i, k] + d[k, j]
d[i, j] = d_ikj if d[i, j] > d_ikj
end
end
end
全ての [i, j] の組は独立して計算できるので、内側の2重ループを map で書き直してみる。
n.times do |k|
d = d.map_with_index do |d_ij, i, j|
d_ikj = d[i, k] + d[k, j]
d_ikj < d_ij ? d_ikj : d_ij
end
end
ループ内の処理をもっと単純にするため、一旦 d_ikj の計算と最小値の計算を分離してみる。（一時記憶用の配列を用意する必要が生じるため、メモリ効率は悪くなる）
n.times do |k|
# （注意）計算時は d にinplaceフラグが立っていないこと
d_via_k = d.map_with_index do |_, i, j|
d[i, k] + d[k, j]
end
d = d.map_with_index do |d_ij, i, j|
d_ikj = d_via_k[i, j]
d_ikj < d_ij ? d_ikj : d_ij
end
end
すると、前半は外積の計算方法と同じもの（演算は掛け算でなく足し算）、後半は要素毎の最小値計算であることがわかる。どちらもNArrayではループを書かず計算できるので、冒頭のコードの形に纏められる。
n.times do |k|
d_via_k = d[true, k..k] + d[k..k, true]
d = Numo::Int64.minimum(d, d_via_k)
end
※ d[true, k..k] と範囲指定しているのは、配列の次元を減らさずに1列を取り出すため。 d.slice(true, k) と書いてもいい。