Remove some dictionaries in star and acyclic coloring

src/SparseMatrixColorings.jl

@@ -39,7 +39,8 @@ using SparseArrays:
     nzrange,
     rowvals,
     sparse,
-    spzeros
+    spzeros,
+    triu
 
 include("graph.jl")
 include("forest.jl")

src/coloring.jl

@@ -84,7 +84,7 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing::
     forbidden_colors = zeros(Int, nv)
     first_neighbor = fill((0, 0), nv)  # at first no neighbors have been encountered
     treated = zeros(Int, nv)
-    star = Dict{Tuple{Int,Int},Int}()
+    star = Vector{Int}(undef, ne)
     sizehint!(star, ne)
     hub = Int[]  # one hub for each star, including the trivial ones
     nb_spokes = Int[]  # number of spokes for each star
@@ -107,7 +107,8 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing::
                 for x in neighbors(g, w)
                     (x == v || iszero(color[x])) && continue
                     wx = _sort(w, x)
-                    if x == hub[star[wx]]  # potential Case 2 (which is always false for trivial stars with two vertices, since the associated hub is negative)
+                    index_wx = g.M[wx...]
+                    if x == hub[star[index_wx]]  # potential Case 2 (which is always false for trivial stars with two vertices, since the associated hub is negative)
                         forbidden_colors[color[x]] = v
                     end
                 end
@@ -121,7 +122,7 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing::
         end
         _update_stars!(star, hub, nb_spokes, g, v, color, first_neighbor)
     end
-    star_set = StarSet(star, hub, nb_spokes)
+    star_set = StarSet(g, star, hub, nb_spokes)
     if postprocessing
         # Reuse the vector forbidden_colors to compute offsets during post-processing
         postprocess!(color, star_set, g, forbidden_colors)
@@ -140,33 +141,44 @@ $TYPEDFIELDS
 """
 struct StarSet
     "a mapping from edges (pair of vertices) to their star index"
-    star::Dict{Tuple{Int,Int},Int}
+    star::Vector{Int}
     "a mapping from star indices to their hub (undefined hubs for single-edge stars are the negative value of one of the vertices, picked arbitrarily)"
     hub::Vector{Int}
     "a mapping from star indices to the vector of their spokes"
     spokes::Vector{Vector{Int}}
+    M::SparseMatrixCSC{Int,Int}
 end
 
-function StarSet(star::Dict{Tuple{Int,Int},Int}, hub::Vector{Int}, nb_spokes::Vector{Int})
+function StarSet(
+    g::AdjacencyGraph{Int}, star::Vector{Int}, hub::Vector{Int}, nb_spokes::Vector{Int}
+)
     # Create a list of spokes for each star, preallocating their sizes based on nb_spokes
     spokes = [Vector{Int}(undef, ns) for ns in nb_spokes]
 
     # Reuse nb_spokes as counters to track the current index while filling the spokes
     fill!(nb_spokes, 0)
 
-    for ((i, j), s) in pairs(star)
-        h = abs(hub[s])
-        nb_spokes[s] += 1
-        index = nb_spokes[s]
-
-        # Assign the non-hub vertex (spoke) to the correct position in spokes
-        if i == h
-            spokes[s][index] = j
-        elseif j == h
-            spokes[s][index] = i
+    nv = nb_vertices(g)
+    k = 0
+    rows = rowvals(g.M)
+    for j in 1:nv
+        for p in nzrange(g.M, j)
+            i = rows[p]
+            k += 1
+            s = star[k]
+            h = abs(hub[s])
+            nb_spokes[s] += 1
+            index = nb_spokes[s]
+
+            # Assign the non-hub vertex (spoke) to the correct position in spokes
+            if i == h
+                spokes[s][index] = j
+            elseif j == h
+                spokes[s][index] = i
+            end
         end
     end
-    return StarSet(star, hub, spokes)
+    return StarSet(star, hub, spokes, g.M)
 end
 
 _sort(u, v) = (min(u, v), max(u, v))
@@ -191,7 +203,7 @@ end
 
 function _update_stars!(
     # modified
-    star::Dict{<:Tuple,<:Integer},
+    star::AbstractVector{<:Integer},
     hub::AbstractVector{<:Integer},
     nb_spokes::AbstractVector{<:Integer},
     # not modified
@@ -203,14 +215,16 @@ function _update_stars!(
     for w in neighbors(g, v)
         iszero(color[w]) && continue
         vw = _sort(v, w)
+        index_vw = g.M[vw...]
         x_exists = false
         for x in neighbors(g, w)
             if x != v && color[x] == color[v]  # vw, wx ∈ E
                 wx = _sort(w, x)
-                star_wx = star[wx]
+                index_wx = g.M[wx...]
+                star_wx = star[index_wx]
                 hub[star_wx] = w  # this may already be true
                 nb_spokes[star_wx] += 1
-                star[vw] = star_wx
+                star[index_vw] = star_wx
                 x_exists = true
                 break
             end
@@ -219,14 +233,15 @@ function _update_stars!(
             (p, q) = first_neighbor[color[w]]
             if p == v && q != w  # vw, vq ∈ E and color[w] = color[q]
                 vq = _sort(v, q)
-                star_vq = star[vq]
+                index_vq = g.M[vq...]
+                star_vq = star[index_vq]
                 hub[star_vq] = v  # this may already be true
                 nb_spokes[star_vq] += 1
-                star[vw] = star_vq
+                star[index_vw] = star_vq
             else  # vw forms a new star
                 push!(hub, -max(v, w))  # star is trivial (composed only of two vertices) so we set the hub to a negative value, but it allows us to choose one of the two vertices
                 push!(nb_spokes, 1)
-                star[vw] = length(hub)
+                star[index_vw] = length(hub)
             end
         end
     end
@@ -251,7 +266,7 @@ This function corresponds to algorithm `DirectRecover2` in the paper.
 function symmetric_coefficient(
     i::Integer, j::Integer, color::AbstractVector{<:Integer}, star_set::StarSet
 )
-    (; star, hub) = star_set
+    (; M, star, hub) = star_set
     if i == j
         # diagonal
         return i, color[j]
@@ -260,7 +275,8 @@ function symmetric_coefficient(
         # star only contains one triangle
         i, j = j, i
     end
-    star_id = star[i, j]
+    index_ij = M[i, j]
+    star_id = star[index_ij]
     h = abs(hub[star_id])
     if h == j
         # i is the spoke
@@ -316,7 +332,7 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessin
                 iszero(color[x]) && continue
                 if forbidden_colors[color[x]] != v
                     _prevent_cycle!(
-                        v, w, x, color, first_visit_to_tree, forbidden_colors, forest
+                        v, w, x, color, g, first_visit_to_tree, forbidden_colors, forest
                     )
                 end
             end
@@ -329,20 +345,20 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessin
         end
         for w in neighbors(g, v)  # grow two-colored stars around the vertex v
             iszero(color[w]) && continue
-            _grow_star!(v, w, color, first_neighbor, forest)
+            _grow_star!(v, w, color, g, first_neighbor, forest)
         end
         for w in neighbors(g, v)
             iszero(color[w]) && continue
             for x in neighbors(g, w)
                 (x == v || iszero(color[x])) && continue
                 if color[x] == color[v]
-                    _merge_trees!(v, w, x, forest)  # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂
+                    _merge_trees!(v, w, x, g, forest)  # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂
                 end
             end
         end
     end
 
-    tree_set = TreeSet(forest, nb_vertices(g))
+    tree_set = TreeSet(g, forest, nb_vertices(g))
     if postprocessing
         # Reuse the vector forbidden_colors to compute offsets during post-processing
         postprocess!(color, tree_set, g, forbidden_colors)
@@ -357,12 +373,14 @@ function _prevent_cycle!(
     x::Integer,
     color::AbstractVector{<:Integer},
     # modified
+    g::AdjacencyGraph,
     first_visit_to_tree::AbstractVector{<:Tuple},
     forbidden_colors::AbstractVector{<:Integer},
     forest::Forest{<:Integer},
 )
     wx = _sort(w, x)
-    id = find_root!(forest, wx)  # The edge wx belongs to the 2-colored tree T, represented by an edge with an integer ID
+    index_wx = g.M[wx...]
+    id = find_root!(forest, index_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)
@@ -379,19 +397,22 @@ function _grow_star!(
     w::Integer,
     color::AbstractVector{<:Integer},
     # modified
+    g::AdjacencyGraph,
     first_neighbor::AbstractVector{<:Tuple},
     forest::Forest{<:Integer},
 )
     vw = _sort(v, w)
-    push!(forest, vw)  # Create a new tree T_{vw} consisting only of edge vw
+    # Create a new tree T_{vw} consisting only of edge vw
     (p, q) = first_neighbor[color[w]]
     if p != v  # a neighbor of v with color[w] encountered for the first time
         first_neighbor[color[w]] = (v, w)
     else  # merge T_{vw} with a two-colored star being grown around v
         vw = _sort(v, w)
         pq = _sort(p, q)
-        root1 = find_root!(forest, vw)
-        root2 = find_root!(forest, pq)
+        index_vw = g.M[vw...]
+        index_pq = g.M[pq...]
+        root1 = find_root!(forest, index_vw)
+        root2 = find_root!(forest, index_pq)
         root_union!(forest, root1, root2)
     end
     return nothing
@@ -403,12 +424,15 @@ function _merge_trees!(
     w::Integer,
     x::Integer,
     # modified
+    g::AdjacencyGraph,
     forest::Forest{<:Integer},
 )
     vw = _sort(v, w)
     wx = _sort(w, x)
-    root1 = find_root!(forest, vw)
-    root2 = find_root!(forest, wx)
+    index_vw = g.M[vw...]
+    index_wx = g.M[wx...]
+    root1 = find_root!(forest, index_vw)
+    root2 = find_root!(forest, index_wx)
     if root1 != root2
         root_union!(forest, root1, root2)
     end
@@ -429,10 +453,8 @@ struct TreeSet
     is_star::Vector{Bool}
 end
 
-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.num_trees: the number of trees in the forest
+function TreeSet(g::AdjacencyGraph{Int}, forest::Forest{Int}, nvertices::Int)
+    # he number of trees in the forest
     nt = forest.num_trees
 
     # dictionary that maps a tree's root to the index of the tree
@@ -444,31 +466,33 @@ function TreeSet(forest::Forest{Int}, nvertices::Int)
 
     # counter of the number of roots found
     k = 0
-    for edge in keys(forest.intmap)
-        i, j = edge
-        root = find_root!(forest, edge)
-
-        # Update roots
-        if !haskey(roots, root)
-            k += 1
-            roots[root] = k
-        end
+    for j in axes(g.M, 2)
+        for edge_index in nzrange(g.M, j)
+            i = g.M.rowval[edge_index]
+            root = find_root!(forest, edge_index)
+
+            # Update roots
+            if !haskey(roots, root)
+                k += 1
+                roots[root] = k
+            end
 
-        # index of the tree T that contains this edge
-        index_tree = roots[root]
+            # index of the tree T that contains this edge
+            index_tree = roots[root]
 
-        # Update the neighbors of i in the tree T
-        if !haskey(trees[index_tree], i)
-            trees[index_tree][i] = [j]
-        else
-            push!(trees[index_tree][i], j)
-        end
+            # Update the neighbors of i in the tree T
+            if !haskey(trees[index_tree], i)
+                trees[index_tree][i] = [j]
+            else
+                push!(trees[index_tree][i], j)
+            end
 
-        # Update the neighbors of j in the tree T
-        if !haskey(trees[index_tree], j)
-            trees[index_tree][j] = [i]
-        else
-            push!(trees[index_tree][j], i)
+            # Update the neighbors of j in the tree T
+            if !haskey(trees[index_tree], j)
+                trees[index_tree][j] = [i]
+            else
+                push!(trees[index_tree][j], i)
+            end
         end
     end
 

src/forest.jl

@@ -10,33 +10,19 @@ Structure that provides fast union-find operations for constructing a forest dur
 $TYPEDFIELDS
 """
 mutable struct Forest{T<:Integer}
-    "current number of edges in the forest"
-    num_edges::T
     "current number of distinct trees in the forest"
     num_trees::T
-    "dictionary mapping each edge represented as a tuple of vertices to its unique integer index"
-    intmap::Dict{Tuple{T,T},T}
     "vector storing the index of a parent in the tree for each edge, used in union-find operations"
     parents::Vector{T}
     "vector approximating the depth of each tree to optimize path compression"
     ranks::Vector{T}
 end
 
 function Forest{T}(n::Integer) where {T<:Integer}
-    num_edges = zero(T)
-    num_trees = zero(T)
-    intmap = Dict{Tuple{T,T},T}()
-    sizehint!(intmap, n)
+    num_trees = T(n)
     parents = collect(Base.OneTo(T(n)))
     ranks = zeros(T, T(n))
-    return Forest{T}(num_edges, num_trees, intmap, parents, ranks)
-end
-
-function Base.push!(forest::Forest{T}, edge::Tuple{T,T}) where {T<:Integer}
-    forest.num_edges += 1
-    forest.intmap[edge] = forest.num_edges
-    forest.num_trees += one(T)
-    return forest
+    return Forest{T}(num_trees, parents, ranks)
 end
 
 function _find_root!(parents::Vector{T}, index_edge::T) where {T<:Integer}
@@ -47,8 +33,8 @@ function _find_root!(parents::Vector{T}, index_edge::T) where {T<:Integer}
     return p
 end
 
-function find_root!(forest::Forest{T}, edge::Tuple{T,T}) where {T<:Integer}
-    return _find_root!(forest.parents, forest.intmap[edge])
+function find_root!(forest::Forest{T}, index_edge::T) where {T<:Integer}
+    return _find_root!(forest.parents, index_edge)
 end
 
 function root_union!(forest::Forest{T}, index_edge1::T, index_edge2::T) where {T<:Integer}