Ruby で解く AtCoder ABC012 D Numo::NArray ワーシャル–フロイド法

はじめに

AtCoder さん、ありがとうございます。

D - バスと避けられない運命

ワーシャル–フロイド法は、最短経路問題で利用されます。

Ruby Numo::NArray なし

n, m = gets.split.map(&:to_i)
ab = Array.new(n) { Array.new(n, 1000000000) }
n.times do |i|
  ab[i][i] = 0
end
m.times do
  a, b, t = gets.split.map(&:to_i)
  ab[a - 1][b - 1] = t
  ab[b - 1][a - 1] = t
end
n.times do |k|
  n.times do |i|
    n.times do |j|
      ab[i][j] = ab[i][k] + ab[k][j] if ab[i][j] > ab[i][k] + ab[k][j]
    end
  end
end
puts ab.map { |e| e.max }.min

他のアルゴリズムに比べ、実装が軽くて助かります。

Ruby Numo::NArray あり

require 'numo/narray'

n, m = gets.split.map(&:to_i)
ab = Numo::Int64.new(n, n).fill(Numo::Int64::MAX >> 1)
ab.diagonal[true] = 0
m.times do
  a, b, t = gets.split.map(&:to_i)
  ab[a - 1, b - 1] = t
  ab[b - 1, a - 1] = t
end
n.times do |k|
	ab = Numo::Int64.minimum(ab, ab[true, k..k] + ab[k..k, true])
end
puts ab.max(axis: 1).min

こちらの記事を参照しました。

いやー、Numo::NArray 速いです。

1/31 追記

コメント欄にてNArrayのmaxメソッドの使い方を教えていただきました。

Pypy(7.3.0)

n, m = map(int, input().split())
ab = [[1000000000] * n for _ in range(n)]
for i in range(n):
    ab[i][i] = 0
for _ in range(m):
    a, b, t = map(int, input().split())
    ab[a - 1][b - 1] = t
    ab[b - 1][a - 1] = t
for k in range(n):
    for i in range(n):
        for j in range(n):
            ab[i][j] = min(ab[i][j], ab[i][k] + ab[k][j])
print(min([max(e) for e in ab]))

Pypy も速いですね。

Crystal(0.33.0)

n, m = read_line.split.map(&.to_i)
ab = Array.new(n) { Array.new(n, 1000000000) }
n.times do |i|
  ab[i][i] = 0
end
m.times do
  a, b, t = read_line.split.map(&.to_i)
  ab[a - 1][b - 1] = t
  ab[b - 1][a - 1] = t
end
n.times do |k|
  n.times do |i|
    n.times do |j|
      ab[i][j] = ab[i][k] + ab[k][j] if ab[i][j] > ab[i][k] + ab[k][j]
    end
  end
end
puts ab.map { |e| e.max }.min

いつものごとく、これを使用しています。

流石に、Crystalが一番速いです。

	Numo::NArray なし	Numo::NArray あり	PyPy3 (7.3.0)	Crystal (0.33.0)
実行時間(ms)	2088	243	523	90

まとめ

• ワーシャル–フロイド法 を学習しました
• Numo::NArray 便利だなと思いました