Acyclic coloring -- enhanced edition

Author: amontoison

src/coloring.jl

@@ -83,14 +83,12 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing::
     ne = nb_edges(g)
     color = zeros(Int, nv)
     forbidden_colors = zeros(Int, nv)
+    edge_to_index = Vector{Int}(undef, nnz(S))
     first_neighbor = fill((0, 0, 0), nv)  # at first no neighbors have been encountered
     treated = zeros(Int, nv)
-    edge_to_index = Vector{Int}(undef, nnz(S))
     star = Vector{Int}(undef, ne)
     hub = Int[]  # one hub for each star, including the trivial ones
-    sizehint!(hub, ne)
     nb_spokes = Int[]  # number of spokes for each star
-    sizehint!(nb_spokes, ne)
     vertices_in_order = vertices(g, order)
 
     # edge_to_index gives an index for each edge
@@ -130,7 +128,7 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing::
             else
                 first_neighbor[color[w]] = (v, w, index_vw)
                 for (ix, x) in enumerate(neighbors2(g, w))
-                    (x == w || x == v || iszero(color[x])) && continue
+                    (w == x || x == v || iszero(color[x])) && continue
                     index_wx = edge_to_index[S.colptr[w] + ix - 1]
                     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
@@ -190,7 +188,7 @@ function _update_stars!(
         index_vw = edge_to_index[S.colptr[v] + iw - 1]
         x_exists = false
         for (ix, x) in enumerate(neighbors2(g, w))
-            (x == w) && continue
+            (w == x) && continue
             if x != v && color[x] == color[v]  # vw, wx ∈ E
                 index_wx = edge_to_index[S.colptr[w] + ix - 1]
                 star_wx = star[index_wx]
@@ -278,8 +276,6 @@ function StarSet(
     return StarSet(star, hub, spokes)
 end
 
-_sort(u, v) = (min(u, v), max(u, v))
-
 """
     acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing::Bool)
 
@@ -305,27 +301,58 @@ If `postprocessing=true`, some colors might be replaced with `0` (the "neutral"
 """
 function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing::Bool)
     # Initialize data structures
+    S = pattern(g)
     nv = nb_vertices(g)
     ne = nb_edges(g)
     color = zeros(Int, nv)
     forbidden_colors = zeros(Int, nv)
-    first_neighbor = fill((0, 0), nv)  # at first no neighbors have been encountered
+    edge_to_index = Vector{Int}(undef, nnz(S))
+    first_neighbor = fill((0, 0, 0), nv)  # at first no neighbors have been encountered
     first_visit_to_tree = fill((0, 0), ne)
     forest = Forest{Int}(ne)
     vertices_in_order = vertices(g, order)
 
+    # edge_to_index gives an index for each edge
+    # use forbidden_colors (or color) for the offsets of each column
+    offsets = forbidden_colors
+    counter = 0
+    rvS = rowvals(S)
+    for j in axes(S, 2)
+        for k in nzrange(S, j)
+            i = rvS[k]
+            if i > j
+                counter += 1
+                edge_to_index[k] = counter
+                k2 = S.colptr[i] + offsets[i]
+                edge_to_index[k2] = counter
+                offsets[i] += 1
+            end
+        end
+    end
+    fill!(offsets, 0)
+    # Note that we don't need to do that for bicoloring,
+    # we can build that in the same time than the transposed sparsity pattern of A
+
     for v in vertices_in_order
         for w in neighbors(g, v)
             iszero(color[w]) && continue
             forbidden_colors[color[w]] = v
         end
         for w in neighbors(g, v)
             iszero(color[w]) && continue
-            for x in neighbors(g, w)
-                iszero(color[x]) && continue
+            for (ix, x) in enumerate(neighbors2(g, w))
+                (w == x || iszero(color[x])) && continue
                 if forbidden_colors[color[x]] != v
+                    index_wx = edge_to_index[S.colptr[w] + ix - 1]
                     _prevent_cycle!(
-                        v, w, x, color, first_visit_to_tree, forbidden_colors, forest
+                        v,
+                        w,
+                        x,
+                        index_wx,
+                        color,
+                        first_visit_to_tree,
+                        forbidden_colors,
+                        forest,
                     )
                 end
             end
@@ -336,22 +363,25 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessin
                 break
             end
         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)
+        for (iw, w) in enumerate(neighbors2(g, v))  # grow two-colored stars around the vertex v
+            (v == w || iszero(color[w])) && continue
+            index_vw = edge_to_index[S.colptr[v] + iw - 1]
+            _grow_star!(v, w, index_vw, color, 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
+        for (iw, w) in enumerate(neighbors2(g, v))
+            (v == w || iszero(color[w])) && continue
+            for (ix, x) in enumerate(neighbors2(g, w))
+                (w == x || 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₂
+                    index_vw = edge_to_index[S.colptr[v] + iw - 1]
+                    index_wx = edge_to_index[S.colptr[w] + ix - 1]
+                    _merge_trees!(v, w, x, index_vw, index_wx, 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, edge_to_index)
     if postprocessing
         # Reuse the vector forbidden_colors to compute offsets during post-processing
         postprocess!(color, tree_set, g, forbidden_colors)
@@ -364,17 +394,16 @@ function _prevent_cycle!(
     v::Integer,
     w::Integer,
     x::Integer,
+    index_wx::Integer,
     color::AbstractVector{<:Integer},
     # modified
     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
+    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)
         first_visit_to_tree[id] = (v, w)
     elseif q != w  # T is connected to vertex v via at least two edges
         forbidden_colors[color[x]] = v
@@ -386,21 +415,19 @@ function _grow_star!(
     # not modified
     v::Integer,
     w::Integer,
+    index_vw::Integer,
     color::AbstractVector{<:Integer},
     # modified
     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
-    (p, q) = first_neighbor[color[w]]
+    # Create a new tree T_{vw} consisting only of edge vw
+    (p, q, index_pq) = 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)
+        first_neighbor[color[w]] = (v, w, index_vw)
     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)
+        root1 = find_root!(forest, index_vw)
+        root2 = find_root!(forest, index_pq)
         root_union!(forest, root1, root2)
     end
     return nothing
@@ -411,13 +438,13 @@ function _merge_trees!(
     v::Integer,
     w::Integer,
     x::Integer,
+    index_vw::Integer,
+    index_wx::Integer,
     # modified
     forest::Forest{<:Integer},
 )
-    vw = _sort(v, w)
-    wx = _sort(w, x)
-    root1 = find_root!(forest, vw)
-    root2 = find_root!(forest, wx)
+    root1 = find_root!(forest, index_vw)
+    root2 = find_root!(forest, index_wx)
     if root1 != root2
         root_union!(forest, root1, root2)
     end
@@ -438,10 +465,9 @@ 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, forest::Forest{Int}, edge_to_index::Vector{Int})
+    S = pattern(g)
+    nv = nb_vertices(g)
     nt = forest.num_trees
 
     # dictionary that maps a tree's root to the index of the tree
@@ -451,38 +477,45 @@ function TreeSet(forest::Forest{Int}, nvertices::Int)
     # vector of dictionaries where each dictionary stores the neighbors of each vertex in a tree
     trees = [Dict{Int,Vector{Int}}() for i in 1:nt]
 
-    # counter of the number of roots found
-    k = 0
-    for edge in keys(forest.intmap)
-        i, j = edge
-        root = find_root!(forest, edge)
+    # current number of roots found
+    nr = 0
 
-        # Update roots
-        if !haskey(roots, root)
-            k += 1
-            roots[root] = k
-        end
+    rvS = rowvals(S)
+    for j in axes(S, 2)
+        for pos in nzrange(S, j)
+            i = rvS[pos]
+            if i > j
+                index_ij = edge_to_index[pos]
+                root = find_root!(forest, index_ij)
 
-        # index of the tree T that contains this edge
-        index_tree = roots[root]
+                # Update roots
+                if !haskey(roots, root)
+                    nr += 1
+                    roots[root] = nr
+                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
+                # index of the tree T that contains this edge
+                index_tree = roots[root]
 
-        # 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 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)
+                end
+            end
         end
     end
 
     # degrees is a vector of integers that stores the degree of each vertex in a tree
-    degrees = Vector{Int}(undef, nvertices)
+    degrees = Vector{Int}(undef, nv)
 
     # reverse breadth first (BFS) traversal order for each tree in the forest
     reverse_bfs_orders = [Tuple{Int,Int}[] for i in 1:nt]

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}