Repository: JuliaCollections/DataStructures.jl
Pull request #599 (2020-03-30 19:02:28)

type-parameterized IntDisjointSets
This allows the creation of type-parameterized `IntDisjointSets`, and keeps the default as `Int` for compatibility. (It also explicitly creates an `IntDisjointSets{Int}` for `DisjointSets`, which does not further type-specialize.)


Patch:
diff --git a/src/DataStructures.jl b/src/DataStructures.jl
index 8f73447..9289933 100644
--- a/src/DataStructures.jl
+++ b/src/DataStructures.jl
@@ -29,7 +29,7 @@ module DataStructures
     export ClassifiedCollections
     export classified_lists, classified_sets, classified_counters
 
-    export IntDisjointSets, DisjointSets, num_groups, find_root, in_same_set, root_union!
+    export IntDisjointSets, DisjointSets, num_groups, find_root!, in_same_set, root_union!
 
     export FenwickTree, length, inc!, dec!, incdec!, prefixsum
 
diff --git a/src/deprecations.jl b/src/deprecations.jl
index 6740785..a98630a 100644
--- a/src/deprecations.jl
+++ b/src/deprecations.jl
@@ -1,3 +1,4 @@
 @deprecate front(x) first(x)
 @deprecate back(x) last(x)
-@deprecate top(x) first(x)
\ No newline at end of file
+@deprecate top(x) first(x)
+@deprecate find_root find_root!   # 2020-03-31
diff --git a/src/disjoint_set.jl b/src/disjoint_set.jl
index c4cb38e..bd83863 100644
--- a/src/disjoint_set.jl
+++ b/src/disjoint_set.jl
@@ -13,23 +13,27 @@
 #
 ############################################################
 
+_intdisjointsets_bounds_err_msg(T) = "the maximum number of elements in IntDisjointSets{$T} is $(typemax(T))"
+
 """
-    IntDisjointSets(n::Integer)
+    IntDisjointSets{T<:Integer}(n::Integer)
 
 A forest of disjoint sets of integers, which is a data structure
 (also called a union–find data structure or merge–find set)
 that tracks a set of elements partitioned
 into a number of disjoint (non-overlapping) subsets.
 """
-mutable struct IntDisjointSets
-    parents::Vector{Int}
-    ranks::Vector{Int}
-    ngroups::Int
+mutable struct IntDisjointSets{T<:Integer}
+    parents::Vector{T}
+    ranks::Vector{T}
+    ngroups::T
 
     # creates a disjoint set comprised of n singletons
-    IntDisjointSets(n::Integer) = new(collect(1:n), zeros(Int, n), n)
+    # IntDisjointSets(n::T) where {T<:Integer} = new(collect(Base.OneTo(n)), zeros(T, n), n)
 end
 
+IntDisjointSets(n::T) where {T<:Integer} = IntDisjointSets{T}(collect(Base.OneTo(n)), zeros(T, n), n)
+IntDisjointSets{T}(n::Integer) where {T<:Integer} = IntDisjointSets{T}(collect(Base.OneTo(T(n))), zeros(T, T(n)), T(n))
 length(s::IntDisjointSets) = length(s.parents)
 
 """
@@ -38,11 +42,11 @@ length(s::IntDisjointSets) = length(s.parents)
 Get a number of groups.
 """
 num_groups(s::IntDisjointSets) = s.ngroups
-Base.eltype(::Type{IntDisjointSets}) = Int
+Base.eltype(::Type{IntDisjointSets{T}}) where {T<:Integer} = T
 
 # find the root element of the subset that contains x
 # path compression is implemented here
-function find_root_impl!(parents::Array{Int}, x::Integer)
+function find_root_impl!(parents::Vector{T}, x::Integer) where {T<:Integer}
     p = parents[x]
     @inbounds if parents[p] != p
         parents[x] = p = _find_root_impl!(parents, p)
@@ -51,7 +55,7 @@ function find_root_impl!(parents::Array{Int}, x::Integer)
 end
 
 # unsafe version of the above
-function _find_root_impl!(parents::Array{Int}, x::Integer)
+function _find_root_impl!(parents::Vector{T}, x::Integer) where {T<:Integer}
     @inbounds p = parents[x]
     @inbounds if parents[p] != p
         parents[x] = p = _find_root_impl!(parents, p)
@@ -60,27 +64,27 @@ function _find_root_impl!(parents::Array{Int}, x::Integer)
 end
 
 """
-    find_root!(s::IntDisjointSets, x::Integer)
+    find_root!(s::IntDisjointSets{T}, x::T)
 
-Find the root element of the subset that contains an Integer x.
+Find the root element of the subset that contains an member x.
 Path compression is implemented here.
 """
-find_root(s::IntDisjointSets, x::Integer) = find_root_impl!(s.parents, x)
+find_root!(s::IntDisjointSets{T}, x::T) where {T<:Integer} = find_root_impl!(s.parents, x)
 
 """
-    in_same_set(s::IntDisjointSets, x::Integer, y::Integer)
+    in_same_set(s::IntDisjointSets{T}, x::T, y::T)
 
 Returns `true` if `x` and `y` belong to the same subset in `s` and `false` otherwise.
 """
-in_same_set(s::IntDisjointSets, x::Integer, y::Integer) = find_root(s, x) == find_root(s, y)
+in_same_set(s::IntDisjointSets{T}, x::T, y::T) where {T<:Integer} = find_root!(s, x) == find_root!(s, y)
 
 """
-    union!(s::IntDisjointSets, x::Integer, y::Integer)
+    union!(s::IntDisjointSets{T}, x::T, y::T)
 
 Merge the subset containing x and that containing y into one
 and return the root of the new set.
 """
-function union!(s::IntDisjointSets, x::Integer, y::Integer)
+function union!(s::IntDisjointSets{T}, x::T, y::T) where {T<:Integer}
     parents = s.parents
     xroot = find_root_impl!(parents, x)
     yroot = find_root_impl!(parents, y)
@@ -88,13 +92,13 @@ function union!(s::IntDisjointSets, x::Integer, y::Integer)
 end
 
 """
-    root_union!(s::IntDisjointSets, x::Integer, y::Integer)
+    root_union!(s::IntDisjointSets{T}, x::T, y::T)
 
 Form a new set that is the union of the two sets whose root elements are
 x and y and return the root of the new set.
 Assume x ≠ y (unsafe).
 """
-function root_union!(s::IntDisjointSets, x::Integer, y::Integer)
+function root_union!(s::IntDisjointSets{T}, x::T, y::T) where {T<:Integer}
     parents = s.parents
     rks = s.ranks
     @inbounds xrank = rks[x]
@@ -103,24 +107,27 @@ function root_union!(s::IntDisjointSets, x::Integer, y::Integer)
     if xrank < yrank
         x, y = y, x
     elseif xrank == yrank
-        rks[x] += 1
+        rks[x] += one(T)
     end
     @inbounds parents[y] = x
-    s.ngroups -= 1
+    s.ngroups -= one(T)
     return x
 end
 
 """
-    push!(s::IntDisjointSets)
+    push!(s::IntDisjointSets{T})
 
 Make a new subset with an automatically chosen new element x.
-Returns the new element.
+Returns the new element. Throw an `ArgumentError` if the
+capacity of the set would be exceeded.
 """
-function push!(s::IntDisjointSets)
-    x = length(s) + 1
+function push!(s::IntDisjointSets{T}) where {T<:Integer}
+    l = length(s)
+    l < typemax(T) || throw(ArgumentError(_intdisjointsets_bounds_err_msg(T))) 
+    x = l + one(T)
     push!(s.parents, x)
-    push!(s.ranks, 0)
-    s.ngroups += 1
+    push!(s.ranks, zero(T))
+    s.ngroups += one(T)
     return x
 end
 
@@ -129,13 +136,13 @@ end
 
 A forest of disjoint sets of arbitrary value type T.
 
-It is a wrapper of IntDisjointSets, which uses a
+It is a wrapper of IntDisjointSets{Int}, which uses a
 dictionary to map the input value to an internal index.
 """
 mutable struct DisjointSets{T} <: AbstractSet{T}
     intmap::Dict{T,Int}
     revmap::Vector{T}
-    internal::IntDisjointSets
+    internal::IntDisjointSets{Int}
 
     DisjointSets{T}() where T = new{T}(Dict{T,Int}(), Vector{T}(), IntDisjointSets(0))
     function DisjointSets{T}(xs) where T    # xs must be iterable
@@ -184,11 +191,11 @@ function sizehint!(s::DisjointSets, n::Integer)
 end
 
 """
-    find_root{T}(s::DisjointSets{T}, x::T)
+    find_root!{T}(s::DisjointSets{T}, x::T)
 
 Finds the root element of the subset in `s` which has the element `x` as a member.
 """
-find_root(s::DisjointSets{T}, x::T) where {T} = s.revmap[find_root(s.internal, s.intmap[x])]
+find_root!(s::DisjointSets{T}, x::T) where {T} = s.revmap[find_root!(s.internal, s.intmap[x])]
 
 """
     in_same_set(s::DisjointSets{T}, x::T, y::T)


Test patch:
diff --git a/test/test_disjoint_set.jl b/test/test_disjoint_set.jl
index 5576dea..23676ee 100644
--- a/test/test_disjoint_set.jl
+++ b/test/test_disjoint_set.jl
@@ -1,52 +1,69 @@
 @testset "DisjointSet" begin
 
     @testset "IntDisjointSets" begin
-        s = IntDisjointSets(10)
-
-        @testset "basic tests" begin
-            @test length(s) == 10
-            @test eltype(s) == Int
-            @test eltype(typeof(s)) == Int
-            @test num_groups(s) == 10
-
-            for i = 1:10
-                @test find_root(s, i) == i
+        for T in [Int, UInt8, Int8, UInt16, Int16, UInt32, Int32, UInt64]
+            @testset "eltype = $(T)" begin
+                s = IntDisjointSets(T(10))
+                s2 = IntDisjointSets{T}(10)
+
+                @testset "basic tests" begin
+                    @test length(s) == 10
+                    @test length(s2) == 10
+                    @test eltype(s) == T
+                    @test eltype(s2) == T
+                    @test eltype(typeof(s)) == T
+                    @test eltype(typeof(s2)) == T
+                    @test num_groups(s) == T(10)
+                    @test num_groups(s2) == T(10)
+
+                    for i = 1:10
+                        @test find_root!(s, T(i)) == T(i)
+                    end
+                    @test_throws BoundsError find_root!(s, T(11))
+
+                    @test !in_same_set(s, T(2), T(3))
+                end
+
+                @testset "union!" begin
+                    union!(s, T(2), T(3))
+                    @test num_groups(s) == T(9)
+                    @test in_same_set(s, T(2), T(3))
+                    @test find_root!(s, T(3)) == T(2)
+                    union!(s, T(3), T(2))
+                    @test num_groups(s) == T(9)
+                    @test in_same_set(s, T(2), T(3))
+                    @test find_root!(s, T(3)) == T(2)
+                end
+
+                @testset "more tests" begin
+                    # We cannot support arbitrary indexing and still use @inbounds with IntDisjointSets
+                    # (and it's not useful anyway)
+                    @test_throws MethodError push!(s, T(22))
+
+                    @test push!(s) == T(11)
+                    @test num_groups(s) == T(10)
+
+                    @test union!(s, T(8), T(7)) == T(8)
+                    @test union!(s, T(5), T(6)) == T(5)
+                    @test union!(s, T(8), T(5)) == T(8)
+                    @test num_groups(s) == T(7)
+                    @test find_root!(s, T(6)) == T(8)
+                    union!(s, T(2), T(6))
+                    @test find_root!(s, T(2)) == T(8)
+                    root1 = find_root!(s, T(6))
+                    root2 = find_root!(s, T(2))
+                    @test root_union!(s, T(root1), T(root2)) == T(8)
+                    @test union!(s, T(5), T(6)) == T(8)
+                end
             end
-            @test_throws BoundsError find_root(s, 11)
-
-            @test !in_same_set(s, 2, 3)
         end
+    end
 
-        @testset "union!" begin
-            union!(s, 2, 3)
-            @test num_groups(s) == 9
-            @test in_same_set(s, 2, 3)
-            @test find_root(s, 3) == 2
-            union!(s, 3, 2)
-            @test num_groups(s) == 9
-            @test in_same_set(s, 2, 3)
-            @test find_root(s, 3) == 2
-        end
-
-        @testset "more tests" begin
-            # We cannot support arbitrary indexing and still use @inbounds with IntDisjointSets
-            # (and it's not useful anyway)
-            @test_throws MethodError push!(s, 22)
-
-            @test push!(s) == 11
-            @test num_groups(s) == 10
-
-            @test union!(s, 8, 7) == 8
-            @test union!(s, 5, 6) == 5
-            @test union!(s, 8, 5) == 8
-            @test num_groups(s) == 7
-            @test find_root(s, 6) == 8
-            union!(s, 2, 6)
-            @test find_root(s, 2) == 8
-            root1 = find_root(s, 6)
-            root2 = find_root(s, 2)
-            @test root_union!(s, root1, root2) == 8
-            @test union!(s, 5, 6) == 8
+    @testset "IntDisjointSets overflow" begin
+        for T in [UInt8, Int8]
+            s = IntDisjointSets(T(typemax(T)-1))
+            push!(s)
+            @test_throws ArgumentError push!(s)
         end
     end
 
@@ -78,7 +95,7 @@
             @test length(s1) == 10
             @test sizehint!(s1, 100) === s1
 
-            r = [find_root(s, i) for i in 1 : 10]
+            r = [find_root!(s, i) for i in 1 : 10]
             @test isequal(r, collect(1:10))
         end
 
@@ -87,20 +104,20 @@
                 x = 2 * i - 1
                 y = 2 * i
                 union!(s, x, y)
-                @test find_root(s, x) == find_root(s, y)
+                @test find_root!(s, x) == find_root!(s, y)
             end
 
 
-            @test union!(s, 1, 4) == find_root(s, 1)
-            @test union!(s, 3, 5) == find_root(s, 1)
-            @test union!(s, 7, 9) == find_root(s, 7)
+            @test union!(s, 1, 4) == find_root!(s, 1)
+            @test union!(s, 3, 5) == find_root!(s, 1)
+            @test union!(s, 7, 9) == find_root!(s, 7)
 
             @test length(s) == 10
             @test num_groups(s) == 2
         end
 
         @testset "r0" begin
-            r0 = [ find_root(s,i) for i in 1:10 ]
+            r0 = [ find_root!(s,i) for i in 1:10 ]
             # Since this is a DisjointSet (not IntDisjointSet), the root for 17 will be 17, not 11
             push!(s, 17)
 
@@ -108,16 +125,16 @@
             @test num_groups(s) == 3
 
             r0 = [ r0 ; 17]
-            r = [find_root(s, i) for i in [1 : 10; 17] ]
+            r = [find_root!(s, i) for i in [1 : 10; 17] ]
             @test isequal(r, r0)
         end
 
         @testset "root_union!" begin
-            root1 = find_root(s, 7)
-            root2 = find_root(s, 3)
+            root1 = find_root!(s, 7)
+            root2 = find_root!(s, 3)
             @test root1 != root2
             root_union!(s, 7, 3)
-            @test find_root(s, 7) == find_root(s, 3)
+            @test find_root!(s, 7) == find_root!(s, 3)
         end
 
         @testset "Some tests using non-integer disjoint sets" begin
@@ -126,21 +143,21 @@
             @test collect(a) == ["a", "b", "c", "d"]
             union!(a, "a", "b")
             @test in_same_set(a,"a","b")
-            @test find_root(a,"a") == find_root(a,"b")
-            @test find_root(a,"a") in elems
+            @test find_root!(a,"a") == find_root!(a,"b")
+            @test find_root!(a,"a") in elems
             @test !in_same_set(a, "c", "d")
             # union returns new root
-            @test find_root(a,"a") == union!(a,"b","c")
+            @test find_root!(a,"a") == union!(a,"b","c")
             union!(a,"c","d")
             # Now they should be in same set, and a is transitively connected to d
             @test in_same_set(a,"a", "d")
             # Root element should thus be same for all:
-            @test all(find_root(a,first(elems)) .== map(x->find_root(a,x),elems))
+            @test all(find_root!(a,first(elems)) .== map(x->find_root!(a,x),elems))
 
-            #@test_throws KeyError find_root(a,"f")
+            #@test_throws KeyError find_root!(a,"f")
 
             push!(a, "f")
-            @test find_root(a,"a") != find_root(a,"f")
+            @test find_root!(a,"a") != find_root!(a,"f")
         end
     end