ABC157 Dの解法メモ【Ruby】

入水目指して過去挑戦して解けてないAtCoder Problems上でDifficultyが水色の問題を解いていく。

ABC157 D: Friend Suggestions概要

同じ"人"の頂点に対して"友達"による辺と"ブロック"による辺の2つの単純無向グラフがある。
それぞれの頂点に対して"友達"になれる頂点(※)の数を求める問題。
※"友達"になれる頂点: "友達"の辺からなるグラフの連結成分であり、"友達"の辺も"ブロック"の辺も持たない頂点。

解法イメージ

問題文の通りに解くと、各頂点に対して"友達"の辺を辿っていき、それぞれの頂点に対して友達かブロックになっているかどうかをチェックしていくことになるが$ \mathcal{O}(N^2) $程度かかってしまうので間に合わない。

連結なグラフがある時、友達になれる可能性のある頂点は連結なグラフにある頂点の数から自身をのぞいた数になる。そのうち、すでに友達になっている頂点とブロックしている頂点を除外した数が実際に友達になれる頂点の数になる。
よって、入力制限から友達の辺やブロックの辺は重複することがないので、Union-Find木を利用して連結な頂点数を求めて各頂点から伸びている友達の辺とブロックの辺の数を引いた数を計算すると答えが出てくる。
ただし、ブロックは連結ではない頂点との辺になっていることがあるので除外すること。

実装例

# Disjoint Set Union
class DSU
  # https://github.com/universato/ac-library-rb/blob/main/lib/dsu.rbのコピペ
end

n, m, k = gets.split.map(&:to_i)
tree = DSU.new(n + 1)
marr = Array.new(n + 1, 0) # すでに友達になっている頂点の数
karr = Array.new(n + 1, 0) # すでにブロックしている頂点の数

m.times do
  a, b = gets.split.map(&:to_i)
  marr[a] += 1
  marr[b] += 1
  tree.merge(a, b)
end
k.times do
  a, b = gets.split.map(&:to_i)
  # 連結ではない頂点のブロックは除外する
  next if !tree.same?(a, b)

  karr[a] += 1
  karr[b] += 1
end

puts (1..n).map { |i| tree.size(i) - marr[i] - karr[i] - 1 }.join(" ")