disjoint set map in julia

IntDisjointMap.jl

struct IntDisjointMap
    parents::Vector{Int64}
    values::Vector{Any}
    merge_fun
end

IntDisjointMap(merge_fun) = IntDisjointMap(Vector{Int64}[], [], merge_fun)

Base.length(x::IntDisjointMap) = length(x.parents)

function Base.push!(x::IntDisjointMap, v) 
    push!(x.parents, -1)
    push!(x.values, v)
    return Base.length(x)
end

function find_root(x::IntDisjointMap, i::Int64)
    while x.parents[i] >= 0 # choice here. I suspect >= is best
        i = x.parents[i]
    end
    return i
end

Base.getindex(x::IntDisjointMap, i::Int64) = x.values[find_root(x,i)]
Base.setindex!(x::IntDisjointMap, v, i::Int64) = x.values[find_root(x,i)] = v


function Base.union!(x::IntDisjointMap, i::Int64, j::Int64)
    pi = find_root(x, i)
    pj = find_root(x, j)
    if pi != pj
        isize = -x.parents[pi]
        jsize = -x.parents[pj]
        if isize > jsize # swap to make size of i less than j
            pi, pj = pj, pi
            isize, jsize = jsize, isize
        end
        x.parents[pj] -= isize # increase new size of pj
        x.values[pj] = x.merge_fun(x.values[pj], x.values[pi])
        x.values[pi] = nothing # clear out unused storage
        x.parents[pi] = pj
    end
    return pj
end

function normalize!(x::IntDisjointMap)
    for i in length(x)
        pi = find_root(x, i)
        if pi != i
            x.parents[i] = pi
        end
    end
end

# Ok the self reference convenion was nice because find_root could stay simple

# If canonized we don't even need a check here.
#function _find_root_canon(x::IntDisjointSet, i::Int64)
#    return @inbounds x.parents[i]
#end
using Test
G = IntDisjointMap(+)
push!(G, 1)
@test G.parents == [-1]
@test G.values == [1]
push!(G,14)
@test G.parents == [-1, -1]
@test G.values == [1, 14]
@test G[1] == 1
@test G[2] == 14
union!(G,1,2)
@test G.parents == [2,-2]
@test G.values == [nothing, 15]
@test G[1] == 15
@test G[2] == 15
push!(G, 4)
@test find_root(G,1) == 2
@test find_root(G,2) == 2
@test find_root(G,3) == 3
union!(G,1,3)
@test G[3] == 19
@test find_root(G,3) == 2
@test find_root(G,1) == 2
@test G.parents == [2,-3,2]

G[3] = 42
@test G[2] == 42

G = IntDisjointMap(+)
push!(G, 42)
push!(G, 13)
@test G[1] == 42
@test G[2] == 13
union!(G, 1, 2)
@test G[1] == 55
@test G[2] == 55