Remove vertices_by_tree for the acyclic coloring (#179)

* Remove vertices_by_tree for the acyclic coloring

* Remove vertices_by_tree for the acyclic coloring

* Update coloring.jl

* Apply suggestions from code review

---------

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

Author: amontoison

src/coloring.jl

@@ -432,7 +432,6 @@ Encode a set of 2-colored trees resulting from the [`acyclic_coloring`](@ref) al
 $TYPEDFIELDS
 """
 struct TreeSet
-    vertices_by_tree::Vector{Vector{Int}}
     reverse_bfs_orders::Vector{Vector{Tuple{Int,Int}}}
 end
 
@@ -445,6 +444,7 @@ function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
 
     # dictionary that maps a tree's root to the index of the tree
     roots = Dict{Int,Int}()
+    sizehint!(roots, ntrees)
 
     # vector of dictionaries where each dictionary stores the neighbors of each vertex in a tree
     trees = [Dict{Int,Vector{Int}}() for i in 1:ntrees]
@@ -454,7 +454,6 @@ function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
     for edge in forest.revmap
         i, j = edge
         # forest has already been compressed so this doesn't change its state
-        # I wanted to use find_root but it is deprecated
         root_edge = find_root!(forest, edge)
         root = forest.intmap[root_edge]
 
@@ -485,14 +484,15 @@ function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
     # degrees is a vector of integers that stores the degree of each vertex in a tree
     degrees = Vector{Int}(undef, nvertices)
 
-    # list of vertices for each tree in the forest
-    vertices_by_tree = [collect(keys(trees[i])) for i in 1:ntrees]
-
     # reverse breadth first (BFS) traversal order for each tree in the forest
     reverse_bfs_orders = [Tuple{Int,Int}[] for i in 1:ntrees]
 
     # nvmax is the number of vertices of the biggest tree in the forest
-    nvmax = mapreduce(length, max, vertices_by_tree; init=0)
+    nvmax = 0
+    for k in 1:ntrees
+        nb_vertices_tree = length(trees[k])
+        nvmax = max(nvmax, nb_vertices_tree)
+    end
 
     # Create a queue with a fixed size nvmax
     queue = Vector{Int}(undef, nvmax)
@@ -542,7 +542,7 @@ function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
         end
     end
 
-    return TreeSet(vertices_by_tree, reverse_bfs_orders)
+    return TreeSet(reverse_bfs_orders)
 end
 
 ## Postprocessing, mirrors decompression code
@@ -554,7 +554,8 @@ function postprocess!(
 )
     (; S) = g
     # flag which colors are actually used during decompression
-    color_used = zeros(Bool, maximum(color))
+    nb_colors = maximum(color)
+    color_used = zeros(Bool, nb_colors)
 
     # nonzero diagonal coefficients force the use of their respective color (there can be no neutral colors if the diagonal is fully nonzero)
     if has_diagonal(g)

src/decompression.jl

@@ -517,7 +517,7 @@ end
 function decompress!(
     A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
 )
-    (; ag, color, vertices_by_tree, reverse_bfs_orders, buffer) = result
+    (; ag, color, reverse_bfs_orders, buffer) = result
     (; S) = ag
     uplo == :F && check_same_pattern(A, S)
     R = eltype(A)
@@ -539,10 +539,14 @@ function decompress!(
     end
 
     # Recover the off-diagonal coefficients of A
-    for k in eachindex(vertices_by_tree, reverse_bfs_orders)
-        for vertex in vertices_by_tree[k]
+    for k in eachindex(reverse_bfs_orders)
+        # Reset the buffer to zero for all vertices in a tree (except the root)
+        for (vertex, _) in reverse_bfs_orders[k]
             buffer_right_type[vertex] = zero(R)
         end
+        # Reset the buffer to zero for the root vertex
+        (_, root) = reverse_bfs_orders[k][end]
+        buffer_right_type[root] = zero(R)
 
         for (i, j) in reverse_bfs_orders[k]
             val = B[i, color[j]] - buffer_right_type[i]
@@ -567,7 +571,6 @@ function decompress!(
     (;
         ag,
         color,
-        vertices_by_tree,
         reverse_bfs_orders,
         diagonal_indices,
         diagonal_nzind,
@@ -612,10 +615,14 @@ function decompress!(
     counter = 0
 
     # Recover the off-diagonal coefficients of A
-    for k in eachindex(vertices_by_tree, reverse_bfs_orders)
-        for vertex in vertices_by_tree[k]
+    for k in eachindex(reverse_bfs_orders)
+        # Reset the buffer to zero for all vertices in a tree (except the root)
+        for (vertex, _) in reverse_bfs_orders[k]
             buffer_right_type[vertex] = zero(R)
         end
+        # Reset the buffer to zero for the root vertex
+        (_, root) = reverse_bfs_orders[k][end]
+        buffer_right_type[root] = zero(R)
 
         for (i, j) in reverse_bfs_orders[k]
             counter += 1

src/result.jl

@@ -278,7 +278,6 @@ struct TreeSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V,R} <:
     ag::G
     color::Vector{Int}
     group::V
-    vertices_by_tree::Vector{Vector{Int}}
     reverse_bfs_orders::Vector{Vector{Tuple{Int,Int}}}
     diagonal_indices::Vector{Int}
     diagonal_nzind::Vector{Int}
@@ -294,7 +293,7 @@ function TreeSetColoringResult(
     tree_set::TreeSet,
     decompression_eltype::Type{R},
 ) where {R}
-    (; vertices_by_tree, reverse_bfs_orders) = tree_set
+    (; reverse_bfs_orders) = tree_set
     (; S) = ag
     nvertices = length(color)
     group = group_by_color(color)
@@ -369,7 +368,6 @@ function TreeSetColoringResult(
         ag,
         color,
         group,
-        vertices_by_tree,
         reverse_bfs_orders,
         diagonal_indices,
         diagonal_nzind,