35 lines
689 B
Ruby
35 lines
689 B
Ruby
|
# Union-find algorithm (weighted quick-union implementation),
|
||
|
# the worst case order of growth is log N,
|
||
|
# see examples in test/union_find_test.rb
|
||
|
class UnionFind
|
||
|
attr_accessor :ids
|
||
|
|
||
|
def initialize(count)
|
||
|
@ids = [*0..count-1]
|
||
|
@sizes = Array.new(count, 1)
|
||
|
end
|
||
|
|
||
|
def connected?(id1, id2)
|
||
|
find(id1) == find(id2)
|
||
|
end
|
||
|
|
||
|
def find(id)
|
||
|
id = @ids[id] while id != @ids[id]
|
||
|
id
|
||
|
end
|
||
|
|
||
|
def union(id1, id2)
|
||
|
root1 = find(id1)
|
||
|
root2 = find(id2)
|
||
|
return if root1 == root2
|
||
|
|
||
|
if @sizes[root1] < @sizes[root2]
|
||
|
@ids[root1] = root2
|
||
|
@sizes[root2] += @sizes[root1]
|
||
|
else
|
||
|
@ids[root2] = root1
|
||
|
@sizes[root1] += @sizes[root2]
|
||
|
end
|
||
|
end
|
||
|
end
|