AtCoder BC097 D: Equals - 銀の弾丸を探して

問題

https://abc097.contest.atcoder.jp/tasks/arc097_b

解法

最初、貪欲に操作をシミュレートして解こうと思ったのですが、組み合わせが果てしなくあることに気づき早々に断念。

例題の解き方を眺めているうちに、各操作のグルーピングがポイントであることに気づきました。

例えば、例題2の場合、

(1,2)と(2,3)は2で繋がってます。このような場合、p₁、p₂、p₃については何回か操作を繰り返すことで任意の並び順に変えることが出来ます。

(x_i, y_i)のグルーピングについては、Union Findのアルゴリズムを使うことで求められます。

グルーピングを行った後、p_iがiと同じグループに入っているればp_i=iとすることが出来るので、その数を数えることで解けます。

実装

class Node
    attr_accessor :parent, :rank
  
    def initialize(n)
      @parent = n
      @rank = 0
    end
  end
  
class UnionFindTree
  def initialize(n)
    @nodes = (0..n).to_a.map { |i| Node.new(i) }
  end

  def find(x)
    return x if @nodes[x].parent == x

    return @nodes[x].parent = find(@nodes[x].parent)
  end

  def unite(a, b)
    a = find(a)
    b = find(b)
    return if a == b

    if @nodes[a].rank < @nodes[b].rank
      @nodes[a].parent = b
    else
      @nodes[b].parent = a
      @nodes[a].rank += 1 if @nodes[a].rank == @nodes[b].rank
    end
  end
 
  def same?(a, b)
    find(a) == find(b)
  end
  
  def parents
    @nodes.map(&:parent)
  end
end
  
n,m = gets.chomp.split.map(&:to_i)
p = gets.chomp.split.map(&:to_i)
tree = UnionFindTree.new(n)
m.times do
  x, y = gets.chomp.split.map(&:to_i)
  tree.unite(x,y)
end

ans = 0
for i in 1..n
    ans += 1 if tree.same?(i, p[i-1])
end
puts ans

Union Findの実装については、以下を参考にさせていただきました。

[Ruby] UnionFind木の実装

こういう定番のアルゴリズムって、事前に用意して置いた方が良いですね。