Union Find Algorithm Implementation in Ruby

require 'pry'

class UnionFind

  attr_accessor :nodes, :data_pairs

  def initialize(n=10, data_set="3-6 4-5")
    @nodes = [] (0..n-1).each { |i| @nodes[i] = i }
    @data_pairs = DataPairs.new(data_set)
  end

  def root(p)
    until nodes[p]==p
      p = nodes[p]
    end

    nodes[p]
  end

  def connected?(p, q)
    root(p) == root(q)
  end

  def build_union
    raise 'No data pairs available' if data_pairs.nil? || data_pairs.data_set.nil?
    data_pairs.data_set.each do |pair|
      union(pair[0], pair[1])
     end
  end

end

class QuickUnionFind < UnionFind

  def union(p, q)
    return if connected?(p, q)
    pid = nodes[p].to_i
    qid = nodes[q].to_i

    (0..nodes.length-1).each do |i|
      nodes[i] = qid if nodes[i] == pid
    end
  end
end

class WeightedQuickUnionFind < UnionFind

  attr_accessor :set_size

  def initialize(n=10, data_set="3-6 4-5")
    @set_size = [1]*n
    super(n, data_set)
  end

  def union(p, q)
    return if connected?(p, q)
    i = root(p)
    j = root(q)
    return if i==j

  if (set_size[i] < set_size[j])
    nodes[i] = j
    set_size[j] += set_size[i] else
    nodes[j] = i
    set_size[i] += set_size[j] end
  end
end

class DataPairs
  attr_accessor :data_set

  def initialize(set="3-6 4-5&#8243;)
    @data_set ||= [] @data_set = make_pairs(set)
  end

  def make_pairs(set)
    number_sets = set.split(" ")
    number_sets.each do |number_set|
      data_set << number_set.split("-").map { |x| x.to_i }
    end
    data_set
  end
end

quf = QuickUnionFind.new(10, "3-6 4-5 9-4 7-0 9-0 3-7")

#quf = QuickUnionFind.new(10, "2-7 9-0 4-8 5-4 5-9 2-0")

quf.build_union

puts "Quick Union: #{quf.nodes.join(" ")}"

wquf = WeightedQuickUnionFind.new(10, "7-2 6-9 5-8 1-3 2-6 8-1 9-4 1-4 3-0")

wquf.build_union

puts "Weighted Quick Union: #{wquf.nodes.join(" ")}"