ac-library-rb で解く AtCoder UnionFind(DSU)

はじめに

AtCoder Problems の Recommendation を利用して、過去の問題を解いています。
AtCoder さん、AtCoder Problems さん、ありがとうございます。

今回のお題

AtCoder Beginner Contest C - Connect Cities
Difficulty: 363

AtCoder Beginner Contest D - Friends
Difficulty: 676

今回のテーマ、UnionFind

ac-library-rb は、AtCoder Library (ACL)のRuby版です。

今回はその中のUnionFind(DSU)を使用しています。

C - Connect Cities

ruby.rb

  
# Disjoint Set Union
class DSU
  def initialize(n = 0)
    # root node: -1 * component size
    # otherwise: parent
    @parent_or_size = Array.new(n, -1)
  end

  def merge(a, b)
    x = leader(a)
    y = leader(b)
    return x if x == y

    x, y = y, x if -@parent_or_size[x] < -@parent_or_size[y]
    @parent_or_size[x] += @parent_or_size[y]
    @parent_or_size[y] = x
    x
  end
  alias unite merge

  def same(a, b)
    leader(a) == leader(b)
  end
  alias same? same

  def leader(a)
    @parent_or_size[a] < 0 ? a : (@parent_or_size[a] = leader(@parent_or_size[a]))
  end
  alias root leader
  alias find leader

  def size(a)
    -@parent_or_size[leader(a)]
  end

  def groups_with_leader
    h = Hash.new { |hash, key| hash[key] = [] }
    @parent_or_size.size.times do |i|
      h[leader(i)] << i
    end
    h
  end

  def groups
    groups_with_leader.values
  end
end

n, m = gets.split.map(&:to_i)
u = DSU.new(n)

m.times do
  a, b = gets.split.map(&:to_i)
  u.merge(a - 1, b - 1)
end

puts u.groups.size - 1

D - Friends

ruby.rb

  
# Disjoint Set Union
class DSU
  def initialize(n = 0)
    # root node: -1 * component size
    # otherwise: parent
    @parent_or_size = Array.new(n, -1)
  end

  def merge(a, b)
    x = leader(a)
    y = leader(b)
    return x if x == y

    x, y = y, x if -@parent_or_size[x] < -@parent_or_size[y]
    @parent_or_size[x] += @parent_or_size[y]
    @parent_or_size[y] = x
    x
  end
  alias unite merge

  def same(a, b)
    leader(a) == leader(b)
  end
  alias same? same

  def leader(a)
    @parent_or_size[a] < 0 ? a : (@parent_or_size[a] = leader(@parent_or_size[a]))
  end
  alias root leader
  alias find leader

  def size(a)
    -@parent_or_size[leader(a)]
  end

  def groups_with_leader
    h = Hash.new { |hash, key| hash[key] = [] }
    @parent_or_size.size.times do |i|
      h[leader(i)] << i
    end
    h
  end

  def groups
    groups_with_leader.values
  end
end

n, m = gets.split.map(&:to_i)
u = DSU.new(n)

m.times do
  a, b = gets.split.map(&:to_i)
  u.merge(a - 1, b - 1)
end

puts u.groups.map { _1.size }.max

一番目のソースに比べ、最後の一行だけ修正しています。

まとめ

• UnionFind(DSU) を解いた
• ACL に詳しくなった
• Ruby に詳しくなった

参照したサイト
ac-library-rb - GitHub
Ruby で解く AtCoder ACL Beginner Contest C UnionFind(DSU) - Qiita
Ruby で解く AtCoder ABC177 D UnionFind - Qiita