Improve the computation of TreeSet (#228)

* Improve the computation of TreeSet

* Update a comment

---------

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

Author: amontoison

src/coloring.jl

@@ -291,7 +291,9 @@ function acyclic_coloring(
         end
     end
 
-    tree_set = TreeSet(g, forest)
+    buffer = forbidden_colors
+    reverse_bfs_orders = first_visit_to_tree
+    tree_set = TreeSet(g, forest, buffer, reverse_bfs_orders)
     if postprocessing
         # Reuse the vector forbidden_colors to compute offsets during post-processing
         offsets = forbidden_colors
@@ -372,19 +374,29 @@ Encode a set of 2-colored trees resulting from the [`acyclic_coloring`](@ref) al
 $TYPEDFIELDS
 """
 struct TreeSet{T}
-    reverse_bfs_orders::Vector{Vector{Tuple{T,T}}}
+    reverse_bfs_orders::Vector{Tuple{T,T}}
     is_star::Vector{Bool}
+    num_edges_per_tree::Vector{T}
 end
 
-function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
+function TreeSet(
+    g::AdjacencyGraph{T},
+    forest::Forest{T},
+    buffer::AbstractVector{T},
+    reverse_bfs_orders::Vector{Tuple{T,T}},
+) where {T}
     S = pattern(g)
     edge_to_index = edge_indices(g)
     nv = nb_vertices(g)
     nt = forest.num_trees
 
-    # dictionary that maps a tree's root to the index of the tree
-    roots = Dict{T,T}()
-    sizehint!(roots, nt)
+    # root_to_tree is a vector that maps a tree's root to the index of the tree
+    # We can recycle forest.ranks because we don't need it anymore to merge trees
+    root_to_tree = forest.ranks
+    fill!(root_to_tree, zero(T))
+
+    # Contains the number of edges per tree
+    num_edges_per_tree = zeros(T, nt)
 
     # vector of dictionaries where each dictionary stores the neighbors of each vertex in a tree
     trees = [Dict{T,Vector{T}}() for i in 1:nt]
@@ -401,13 +413,14 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
                 root = find_root!(forest, index_ij)
 
                 # Update roots
-                if !haskey(roots, root)
+                if iszero(root_to_tree[root])
                     nr += 1
-                    roots[root] = nr
+                    root_to_tree[root] = nr
                 end
 
                 # index of the tree T that contains this edge
-                index_tree = roots[root]
+                index_tree = root_to_tree[root]
+                num_edges_per_tree[index_tree] += 1
 
                 # Update the neighbors of i in the tree T
                 if !haskey(trees[index_tree], i)
@@ -427,10 +440,7 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
     end
 
     # degrees is a vector of integers that stores the degree of each vertex in a tree
-    degrees = Vector{T}(undef, nv)
-
-    # reverse breadth first (BFS) traversal order for each tree in the forest
-    reverse_bfs_orders = [Tuple{T,T}[] for i in 1:nt]
+    degrees = buffer
 
     # nvmax is the number of vertices of the biggest tree in the forest
     nvmax = 0
@@ -446,6 +456,10 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
     # meaning that one vertex is directly connected to all other vertices in the tree
     is_star = Vector{Bool}(undef, nt)
 
+    # Number of edges treated
+    num_edges_treated = zero(T)
+
+    # reverse_bfs_orders contains the reverse breadth first (BFS) traversal order for each tree in the forest
     for k in 1:nt
         tree = trees[k]
 
@@ -483,7 +497,8 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
                 # Check if neighbor is the parent of the leaf or if it was a child before the tree was pruned
                 if degrees[neighbor] != 0
                     # (leaf, neighbor) represents the next edge to visit during decompression
-                    push!(reverse_bfs_orders[k], (leaf, neighbor))
+                    num_edges_treated += 1
+                    reverse_bfs_orders[num_edges_treated] = (leaf, neighbor)
 
                     if bool_star
                         # Initialize the potential hub of the star with the first parent of a leaf
@@ -514,7 +529,7 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
         is_star[k] = bool_star
     end
 
-    return TreeSet(reverse_bfs_orders, is_star)
+    return TreeSet(reverse_bfs_orders, is_star, num_edges_per_tree)
 end
 
 ## Postprocessing, mirrors decompression code
@@ -582,46 +597,52 @@ function postprocess!(
         end
     else
         # only the colors of non-leaf vertices are used
-        (; reverse_bfs_orders, is_star) = star_or_tree_set
+        (; reverse_bfs_orders, is_star, num_edges_per_tree) = star_or_tree_set
         nb_trivial_trees = 0
 
+        # Index of the first edge in reverse_bfs_orders for the current tree
+        first = 1
+
         # Iterate through all non-trivial trees
-        for k in eachindex(reverse_bfs_orders)
-            reverse_bfs_order = reverse_bfs_orders[k]
+        for k in eachindex(num_edges_per_tree)
+            ne_tree = num_edges_per_tree[k]
             # Check if we have more than one edge in the tree (non-trivial tree)
-            if length(reverse_bfs_order) > 1
+            if ne_tree > 1
                 # Determine if the tree is a star
                 if is_star[k]
                     # It is a non-trivial star and only the color of the hub is needed
-                    (_, hub) = reverse_bfs_order[1]
+                    (_, hub) = reverse_bfs_orders[first]
                     color_used[color[hub]] = true
                 else
                     # It is not a star and both colors are needed during the decompression
-                    (i, j) = reverse_bfs_order[1]
+                    (i, j) = reverse_bfs_orders[first]
                     color_used[color[i]] = true
                     color_used[color[j]] = true
                 end
             else
                 nb_trivial_trees += 1
             end
+            first += ne_tree
         end
 
         # Process the trivial trees (if any)
         if nb_trivial_trees > 0
-            for k in eachindex(reverse_bfs_orders)
-                reverse_bfs_order = reverse_bfs_orders[k]
+            first = 1
+            for k in eachindex(num_edges_per_tree)
+                ne_tree = num_edges_per_tree[k]
                 # Check if we have exactly one edge in the tree
-                if length(reverse_bfs_order) == 1
-                    (i, j) = reverse_bfs_order[1]
+                if ne_tree == 1
+                    (i, j) = reverse_bfs_orders[first]
                     if color_used[color[i]]
                         # Make i the root to avoid possibly adding one more used color
                         # Switch it with the (only) leaf
-                        reverse_bfs_order[1] = (j, i)
+                        reverse_bfs_orders[first] = (j, i)
                     else
                         # Keep j as the root
                         color_used[color[j]] = true
                     end
                 end
+                first += ne_tree
             end
         end
     end

src/decompression.jl

@@ -517,7 +517,7 @@ end
 function decompress!(
     A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
 )
-    (; ag, color, reverse_bfs_orders, buffer) = result
+    (; ag, color, reverse_bfs_orders, num_edges_per_tree, buffer) = result
     (; S) = ag
     uplo == :F && check_same_pattern(A, S)
     R = eltype(A)
@@ -531,24 +531,32 @@ function decompress!(
 
     # Recover the diagonal coefficients of A
     if has_diagonal(ag)
-        for i in axes(A, 1)
+        for i in axes(S, 1)
             if !iszero(S[i, i])
                 A[i, i] = B[i, color[i]]
             end
         end
     end
 
+    # Index of the first edge in reverse_bfs_orders for the current tree
+    first = 1
+
     # Recover the off-diagonal coefficients of A
-    for k in eachindex(reverse_bfs_orders)
+    for k in eachindex(num_edges_per_tree)
+        ne_tree = num_edges_per_tree[k]
+        last = first + ne_tree - 1
+
         # Reset the buffer to zero for all vertices in a tree (except the root)
-        for (vertex, _) in reverse_bfs_orders[k]
+        for pos in first:last
+            (vertex, _) = reverse_bfs_orders[pos]
             buffer_right_type[vertex] = zero(R)
         end
         # Reset the buffer to zero for the root vertex
-        (_, root) = reverse_bfs_orders[k][end]
+        (_, root) = reverse_bfs_orders[last]
         buffer_right_type[root] = zero(R)
 
-        for (i, j) in reverse_bfs_orders[k]
+        for pos in first:last
+            (i, j) = reverse_bfs_orders[pos]
             val = B[i, color[j]] - buffer_right_type[i]
             buffer_right_type[j] = buffer_right_type[j] + val
 
@@ -559,6 +567,7 @@ function decompress!(
                 A[j, i] = val
             end
         end
+        first += ne_tree
     end
     return A
 end
@@ -573,6 +582,7 @@ function decompress!(
         ag,
         color,
         reverse_bfs_orders,
+        num_edges_per_tree,
         diagonal_indices,
         diagonal_nzind,
         lower_triangle_offsets,
@@ -612,20 +622,28 @@ function decompress!(
         end
     end
 
+    # Index of the first edge in reverse_bfs_orders for the current tree
+    first = 1
+
     # Index of offsets in lower_triangle_offsets and upper_triangle_offsets
     counter = 0
 
     # Recover the off-diagonal coefficients of A
-    for k in eachindex(reverse_bfs_orders)
+    for k in eachindex(num_edges_per_tree)
+        ne_tree = num_edges_per_tree[k]
+        last = first + ne_tree - 1
+
         # Reset the buffer to zero for all vertices in a tree (except the root)
-        for (vertex, _) in reverse_bfs_orders[k]
+        for pos in first:last
+            (vertex, _) = reverse_bfs_orders[pos]
             buffer_right_type[vertex] = zero(R)
         end
         # Reset the buffer to zero for the root vertex
-        (_, root) = reverse_bfs_orders[k][end]
+        (_, root) = reverse_bfs_orders[last]
         buffer_right_type[root] = zero(R)
 
-        for (i, j) in reverse_bfs_orders[k]
+        for pos in first:last
+            (i, j) = reverse_bfs_orders[pos]
             counter += 1
             val = B[i, color[j]] - buffer_right_type[i]
             buffer_right_type[j] = buffer_right_type[j] + val
@@ -665,6 +683,7 @@ function decompress!(
             end
             #! format: on
         end
+        first += ne_tree
     end
     return A
 end

src/result.jl

@@ -314,7 +314,8 @@ struct TreeSetColoringResult{
     ag::G
     color::Vector{T}
     group::GT
-    reverse_bfs_orders::Vector{Vector{Tuple{T,T}}}
+    reverse_bfs_orders::Vector{Tuple{T,T}}
+    num_edges_per_tree::Vector{T}
     diagonal_indices::Vector{T}
     diagonal_nzind::Vector{T}
     lower_triangle_offsets::Vector{T}
@@ -329,7 +330,7 @@ function TreeSetColoringResult(
     tree_set::TreeSet{<:Integer},
     decompression_eltype::Type{R},
 ) where {T<:Integer,R}
-    (; reverse_bfs_orders) = tree_set
+    (; reverse_bfs_orders, num_edges_per_tree) = tree_set
     (; S) = ag
     nvertices = length(color)
     group = group_by_color(T, color)
@@ -358,11 +359,18 @@ function TreeSetColoringResult(
     lower_triangle_offsets = Vector{T}(undef, nedges)
     upper_triangle_offsets = Vector{T}(undef, nedges)
 
+    # Index of the first edge in reverse_bfs_orders for the current tree
+    first = 1
+
     # Index in lower_triangle_offsets and upper_triangle_offsets
     index_offsets = 0
 
-    for k in eachindex(reverse_bfs_orders)
-        for (leaf, neighbor) in reverse_bfs_orders[k]
+    for k in eachindex(num_edges_per_tree)
+        ne_tree = num_edges_per_tree[k]
+        last = first + ne_tree - 1
+
+        for pos in first:last
+            (leaf, neighbor) = reverse_bfs_orders[pos]
             # Update lower_triangle_offsets and upper_triangle_offsets
             i = leaf
             j = neighbor
@@ -393,6 +401,7 @@ function TreeSetColoringResult(
             end
             #! format: on
         end
+        first += ne_tree
     end
 
     # buffer holds the sum of edge values for subtrees in a tree.
@@ -405,6 +414,7 @@ function TreeSetColoringResult(
         color,
         group,
         reverse_bfs_orders,
+        num_edges_per_tree,
         diagonal_indices,
         diagonal_nzind,
         lower_triangle_offsets,