Princeton Algorithms - Week One - Weighted Quick Union

weighted-quick-union.rb

class WeightedQuickUnion
  attr_accessor :id

  def initialize(n)
    @id = Array.new(n) { |e| e = e }
    @size = Array.new(n) { |e| e = 1 }
  end

  def connect(a, b)
    c = parent_of(a)
    d = parent_of(b)

    return if c == d

    if (@size[c] < @size[d])
      @id[c] = d
      @size[d] += @size[c]
    else
      @id[d] = c
      @size[c] += @size[d]
    end
  end

  def connected?(a, b)
    parent_of(a) === parent_of(b)
  end

  private

  def parent_of(n)
    parent = n
    while parent != @id[parent] do
      # @id[parent] = @id[@id[parent]] # path compression (optimization)
      parent = @id[parent]
    end
    parent
  end
end

wqu = WeightedQuickUnion.new(10)

wqu.connect(4, 3)
wqu.connect(3, 2)
wqu.connect(2, 1)
wqu.connect(1, 8)
puts wqu.id.inspect
puts wqu.connected?(4, 8)