diff --git a/src/FunctionalCollections.jl b/src/FunctionalCollections.jl index 5f4a0e7..952b670 100644 --- a/src/FunctionalCollections.jl +++ b/src/FunctionalCollections.jl @@ -46,6 +46,10 @@ include("PersistentQueue.jl") export PersistentQueue, queue, enq +include("PersistentDisjointSet.jl") + +export PersistentDisjointSet + export @Persistent fromexpr(ex::Expr, ::Type{pvec}) = :(pvec($(esc(ex)))) diff --git a/src/PersistentDisjointSet.jl b/src/PersistentDisjointSet.jl new file mode 100644 index 0000000..673fb3c --- /dev/null +++ b/src/PersistentDisjointSet.jl @@ -0,0 +1,32 @@ +struct PersistentDisjointSet{T} + table::PersistentHashMap{T,T} + heights::PersistentHashMap{T,UInt8} +end +PersistentDisjointSet{T}() where {T} = PersistentDisjointSet{T}(PersistentHashMap{T,T}(),PersistentHashMap{T,UInt8}()) + + +height(t::PersistentDisjointSet{T},i::T) where {T} = get(t.heights,i,zero(UInt8)) + +function Base.getindex(t::PersistentDisjointSet{T},i::T) where {T} + while haskey(t.table,i) # In a persistent setting, writes cost memory, so path compression may not be worth the allocations. + i = t.table[i] # Height balancing already guarentees finding the root in floor(lg(N)) lookups. + end + i +end + +function Base.union(t::PersistentDisjointSet{T},i::T,j::T) where {T} + i = t[i];j=t[j] + if i == j return t end + if height(t,i) > height(t,j) + PersistentDisjointSet(assoc(t.table,j,i),t.heights) + elseif height(t,i) < height(t,j) + PersistentDisjointSet(assoc(t.table,i,j),t.heights) + else + PersistentDisjointSet(assoc(t.table,j,i),assoc(t.heights,i,height(t,i) + 1)) + end +end + +function Base.show(io::IO,t::PersistentDisjointSet{T}) where {T} + Base.print(io,"PersistentDisjointSet{$T}") + Base.print(io,sort!(map(x->first(x) => t[last(x)],collect(pairs(t.table)));by=x->(last(x),first(x)))) +end \ No newline at end of file diff --git a/test/PersistentDisjointSetTest.jl b/test/PersistentDisjointSetTest.jl new file mode 100644 index 0000000..4b6abbf --- /dev/null +++ b/test/PersistentDisjointSetTest.jl @@ -0,0 +1,19 @@ +using FunctionalCollections +using Test + +@testset "Persistent Lists" begin + + @testset "UnionFind" begin + uf = PersistentDisjointSet{Int}() + uf = union(uf,2,3) + uf = union(uf,5,7) + uf = union(uf,2,4) + uf = union(uf,4,8) + @assert uf[1] == 1 + @assert uf[1] != uf[7] + @assert uf[3] == uf[8] + @assert uf[5] == uf[7] + @assert uf[3] != uf[7] + end + +end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 80cd5e4..ee12b3d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -4,4 +4,5 @@ include("./PersistentMapTest.jl") include("./PersistentSetTest.jl") include("./PersistentListTest.jl") include("./PersistentQueueTest.jl") -include("./PersistentMacroTest.jl") \ No newline at end of file +include("./PersistentDisjointSetTest.jl") +include("./PersistentMacroTest.jl")