Implement our own structure Forest

Author: amontoison

Project.toml

@@ -5,7 +5,6 @@ version = "0.4.15"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
-DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -23,7 +22,6 @@ SparseMatrixColoringsColorsExt = "Colors"
 ADTypes = "1.2.1"
 CliqueTrees = "0.5.2"
 Colors = "0.12.11, 0.13"
-DataStructures = "0.18"
 DocStringExtensions = "0.8,0.9"
 LinearAlgebra = "<0.0.1, 1"
 Random = "<0.0.1, 1"

src/SparseMatrixColorings.jl

@@ -11,7 +11,6 @@ module SparseMatrixColorings
 
 using ADTypes: ADTypes
 using Base.Iterators: Iterators
-using DataStructures: DisjointSets, find_root!, root_union!, num_groups
 using DocStringExtensions: README, EXPORTS, SIGNATURES, TYPEDEF, TYPEDFIELDS
 using LinearAlgebra:
     Adjoint,
@@ -43,6 +42,7 @@ using SparseArrays:
     spzeros
 
 include("graph.jl")
+include("forest.jl")
 include("order.jl")
 include("coloring.jl")
 include("result.jl")

src/coloring.jl

@@ -302,11 +302,7 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessin
     forbidden_colors = zeros(Int, nv)
     first_neighbor = fill((0, 0), nv)  # at first no neighbors have been encountered
     first_visit_to_tree = fill((0, 0), ne)
-    forest = DisjointSets{Tuple{Int,Int}}()
-    sizehint!(forest.intmap, ne)
-    sizehint!(forest.revmap, ne)
-    sizehint!(forest.internal.parents, ne)
-    sizehint!(forest.internal.ranks, ne)
+    forest = Forest{Int}(ne)
     vertices_in_order = vertices(g, order)
 
     for v in vertices_in_order
@@ -347,7 +343,7 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessin
     end
 
     # compress forest
-    for edge in forest.revmap
+    for edge in keys(forest.intmap)
         find_root!(forest, edge)
     end
     tree_set = TreeSet(forest, nb_vertices(g))
@@ -367,11 +363,10 @@ function _prevent_cycle!(
     # modified
     first_visit_to_tree::AbstractVector{<:Tuple},
     forbidden_colors::AbstractVector{<:Integer},
-    forest::DisjointSets{<:Tuple{Int,Int}},
+    forest::Forest{<:Integer},
 )
     wx = _sort(w, x)
-    root = find_root!(forest, wx)  # edge wx belongs to the 2-colored tree T represented by edge "root"
-    id = forest.intmap[root] # ID of the representative edge "root" of a two-colored tree T.
+    id = find_root!(forest, wx)  # The edge wx belongs to the 2-colored tree T, represented by an edge with an integer ID
     (p, q) = first_visit_to_tree[id]
     if p != v  # T is being visited from vertex v for the first time
         vw = _sort(v, w)
@@ -389,7 +384,7 @@ function _grow_star!(
     color::AbstractVector{<:Integer},
     # modified
     first_neighbor::AbstractVector{<:Tuple},
-    forest::DisjointSets{Tuple{Int,Int}},
+    forest::Forest{<:Integer},
 )
     vw = _sort(v, w)
     push!(forest, vw)  # Create a new tree T_{vw} consisting only of edge vw
@@ -412,7 +407,7 @@ function _merge_trees!(
     w::Integer,
     x::Integer,
     # modified
-    forest::DisjointSets{Tuple{Int,Int}},
+    forest::Forest{<:Integer},
 )
     vw = _sort(v, w)
     wx = _sort(w, x)
@@ -438,12 +433,11 @@ struct TreeSet
     is_star::Vector{Bool}
 end
 
-function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
-    # forest is a structure DisjointSets from DataStructures.jl
+function TreeSet(forest::Forest{Int}, nvertices::Int)
+    # Forest is a structure defined in forest.jl
     # - forest.intmap: a dictionary that maps an edge (i, j) to an integer k
-    # - forest.revmap: a dictionary that does the reverse of intmap, mapping an integer k to an edge (i, j)
-    # - forest.internal.ngroups: the number of trees in the forest
-    ntrees = forest.internal.ngroups
+    # - forest.ntrees: the number of trees in the forest
+    ntrees = forest.ntrees
 
     # dictionary that maps a tree's root to the index of the tree
     roots = Dict{Int,Int}()
@@ -454,11 +448,9 @@ function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
 
     # counter of the number of roots found
     k = 0
-    for edge in forest.revmap
+    for edge in keys(forest.intmap)
         i, j = edge
-        # forest has already been compressed so this doesn't change its state
-        root_edge = find_root!(forest, edge)
-        root = forest.intmap[root_edge]
+        root = find_root!(forest, edge)
 
         # Update roots
         if !haskey(roots, root)

src/forest.jl

@@ -0,0 +1,52 @@
+mutable struct Forest{T<:Integer}
+    counter::T
+    intmap::Dict{Tuple{T,T},T}
+    parents::Vector{T}
+    ranks::Vector{T}
+    ntrees::T
+end
+
+function Forest{T}(n::Integer) where {T<:Integer}
+    counter = zero(T)
+    intmap = Dict{Tuple{T,T},T}()
+    sizehint!(intmap, n)
+    parents = collect(Base.OneTo(T(n)))
+    ranks = zeros(T, T(n))
+    ntrees = T(n)
+    return Forest{T}(counter, intmap, parents, ranks, ntrees)
+end
+
+function Base.push!(forest::Forest{T}, edge::Tuple{T,T}) where {T<:Integer}
+    forest.counter += 1
+    forest.intmap[edge] = forest.counter
+    forest.ntrees += one(T)
+    return edge
+end
+
+function _find_root!(parents::Vector{T}, index_edge::T) where {T<:Integer}
+    @inbounds p = parents[index_edge]
+    @inbounds if parents[p] != p
+        parents[index_edge] = p = _find_root!(parents, p)
+    end
+    return p
+end
+
+function find_root!(forest::Forest{T}, edge::Tuple{T,T}) where {T<:Integer}
+    return _find_root!(forest.parents, forest.intmap[edge])
+end
+
+function root_union!(forest::Forest{T}, index_edge1::T, index_edge2::T) where {T<:Integer}
+    parents = forest.parents
+    rks = forest.ranks
+    @inbounds rank1 = rks[index_edge1]
+    @inbounds rank2 = rks[index_edge2]
+
+    if rank1 < rank2
+        index_edge1, index_edge2 = index_edge2, index_edge1
+    elseif rank1 == rank2
+        rks[index_edge1] += one(T)
+    end
+    @inbounds parents[index_edge2] = index_edge1
+    forest.ntrees -= one(T)
+    return nothing
+end