Ruby で解く AtCoder ABC292 E DFS ワーシャル–フロイド法 + bit演算

はじめに

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

今回は、コンテスト後直ぐに解説放送が始まりました。

その中のE問題を写経したいと思います。

E - Transitivity

DFS

Ruby

n, m = gets.split.map(&:to_i)
to = Array.new(n) { [] }
m.times do
  a, b = gets.split.map(&:to_i)
  to[a - 1] << b - 1
end
ans = 0
n.times do |sv|
  can = Array.new(n, false)
  q = [sv]
  can[sv] = true
  while q.empty?.!
    v = q.pop
    to[v].each do |u|
      next if can[u]
      can[u] = true
      q << u
      ans += 1
    end
  end
end
puts ans - m

いつものごとく、あれを使ってCrystal

Crystal

n, m = read_line.split.map(&.to_i)  
to = Array.new(n) { [] of Int32 }   
m.times do
  a, b = read_line.split.map(&.to_i)
  to[a - 1] << b - 1
end
ans = 0
n.times do |sv|
  can = Array.new(n, false)
  q = [sv]
  can[sv] = true
  while q.empty?.!
    v = q.pop
    to[v].each do |u|
      next if can[u]
      can[u] = true
      q << u
      ans += 1
    end
  end
end
puts ans - m

あれ

ワーシャル–フロイド法(TLE)

n, m = gets.split.map(&:to_i)
ab = Array.new(n) { Array.new(n, 0) }
n.times do |i|
  ab[i][i] = 1
end
m.times do
  a, b = gets.split.map(&:to_i)
  ab[a - 1][b - 1] = 1
end
n.times do |k|
  n.times do |i|
    n.times do |j|
      ab[i][j] |= ab[i][k] & ab[k][j]
    end
  end
end
puts ab.flatten.count(1) - m - n

ワーシャル–フロイド法で距離を全て1としますとDFS/BFSの問題を解くことができます。
但し、DFSの計算量は$O(N^{2})$ですが、ワーシャル–フロイド法は$O(N^{3})$ですので、この問題ではTLEとなります。

ワーシャル–フロイド法 + bit演算

n, m = gets.split.map(&:to_i)
g = Array.new(n, 0)
n.times do |i|
  g[i] |= 2 ** (n - i - 1)
end
m.times do
  a, b = gets.split.map(&:to_i)
  g[a - 1] |= 2 ** (n - b)
end
n.times do |k|
  n.times do |i|
    g[i] |= g[k] if g[i][n - k - 1] == 1
  end
end
puts g.map { |e| e.to_s(2).count('1') }.sum - m - n

bit演算を使用することにより、3重ループを2重ループに落とすことができました。

	DFS(Ruby)	DFS(Crystal)	FW	FW+bitset
実行時間(ms)	561	66	TLE	1036

まとめ

• ワーシャル–フロイド法の応用 を学習しました
• すぬけさんありがとうございます。