DisjointSetUnion(Ruby)

DSU.rb

class DisjointSetUnion
  def initialize(size)
    @nodes = Array.new(size) { |i| TreeNode.new(i) }
  end

  def unite(i, j)
    @nodes[i].unite(@nodes[j])
  end

  def same_group?(i, j)
    @nodes[i].in_same_tree?(@nodes[j])
  end

  def group_size(i)
    @nodes[i].tree_size
  end
end

class TreeNode
  attr_accessor :parent, :descendant_cnt

  def initialize(i)
    @id = i
    @parent = self
    @descendant_cnt = 1
  end

  def root(node = self)
    return node if node.parent == node
    node.parent = root(node.parent)
  end

  def in_same_tree?(other)
    self.root == other.root
  end

  def unite(other)
    root_a, root_b = self.root, other.root
    root_a, root_b = root_b, root_a if root_a.descendant_cnt > root_b.descendant_cnt

    root_a.parent = root_b
    root_b.descendant_cnt += root_a.descendant_cnt
  end

  def tree_size
    self.root.descendant_cnt
  end
end

N, Q = gets.split(' ').map(&:to_i)
union = DisjointSetUnion.new(N)
Q.times do
  t, a, b = gets.split(' ').map(&:to_i)
  if t == 0
    union.unite(a - 1, b - 1)
  else
    puts union.same_group?(a - 1, b - 1) ? 'yes' : 'no'
  end
end