Detect stars in the forest for acyclic coloring

Author: amontoison

src/coloring.jl

@@ -434,6 +434,7 @@ $TYPEDFIELDS
 struct TreeSet
     vertices_by_tree::Vector{Vector{Int}}
     reverse_bfs_orders::Vector{Vector{Tuple{Int,Int}}}
+    is_star::Vector{Bool}
 end
 
 function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
@@ -497,9 +498,19 @@ function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
     # Create a queue with a fixed size nvmax
     queue = Vector{Int}(undef, nvmax)
 
+    # Specify if each tree in the forest is a star,
+    # meaning that one vertex is directly connected to all other vertices in the tree
+    is_star = Vector{Bool}(undef, ntrees)
+
     for k in 1:ntrees
         tree = trees[k]
 
+        # Boolean indicating whether the current tree is a star (a single central vertex connected to all others)
+        bool_star = true
+
+        # Candidate hub vertex if the current tree is a star
+        virtual_hub = 0
+
         # Initialize the queue to store the leaves
         queue_start = 1
         queue_end = 0
@@ -529,6 +540,18 @@ function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
                     # (leaf, neighbor) represents the next edge to visit during decompression
                     push!(reverse_bfs_orders[k], (leaf, neighbor))
 
+                    if bool_star
+                        # Initialize the potential hub of the star with the first parent of a leaf
+                        if virtual_hub == 0
+                            virtual_hub = neighbor
+                        else
+                            # Verify if the tree still qualifies as a star
+                            if virtual_hub != neighbor
+                                bool_star = false
+                            end
+                        end
+                    end
+
                     # reduce the degree of the neighbor
                     degrees[neighbor] -= 1
 
@@ -540,9 +563,12 @@ function TreeSet(forest::DisjointSets{Tuple{Int,Int}}, nvertices::Int)
                 end
             end
         end
+
+        # Specify if the tree is a star or not
+        is_star[k] = bool_star
     end
 
-    return TreeSet(vertices_by_tree, reverse_bfs_orders)
+    return TreeSet(vertices_by_tree, reverse_bfs_orders, is_star)
 end
 
 ## Postprocessing, mirrors decompression code
@@ -599,17 +625,23 @@ function postprocess!(
         end
     else
         # only the colors of non-leaf vertices are used
-        (; reverse_bfs_orders) = star_or_tree_set
+        (; reverse_bfs_orders, is_star) = star_or_tree_set
         nb_trivial_trees = 0
 
         # Iterate through all non-trivial trees
         for k in eachindex(reverse_bfs_orders)
             reverse_bfs_order = reverse_bfs_orders[k]
-            # Check if we have more than one edge in the tree
+            # Check if we have more than one edge in the tree (non-trivial tree)
             if length(reverse_bfs_order) > 1
-                # TODO: Optimize by avoiding iteration over all edges
-                # Only one edge is needed if we know if it is a normal tree or a star
-                for (i, j) in reverse_bfs_order
+                # 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[end]
+                    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]
+                    color_used[color[i]] = true
                     color_used[color[j]] = true
                 end
             else

src/decompression.jl

@@ -505,7 +505,7 @@ end
 function decompress!(
     A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
 )
-    (; color, vertices_by_tree, reverse_bfs_orders, buffer) = result
+    (; color, vertices_by_tree, reverse_bfs_orders, is_star, buffer) = result
     S = result.ag.S
     uplo == :F && check_same_pattern(A, S)
     R = eltype(A)
@@ -526,13 +526,22 @@ function decompress!(
 
     # Recover the off-diagonal coefficients of A
     for k in eachindex(vertices_by_tree, reverse_bfs_orders)
-        for vertex in vertices_by_tree[k]
-            buffer_right_type[vertex] = zero(R)
+        # We need the buffer only when the tree is not a star (trivial or non-trivial)
+        if !is_star[k]
+            for vertex in vertices_by_tree[k]
+                buffer_right_type[vertex] = zero(R)
+            end
         end
 
         for (i, j) in reverse_bfs_orders[k]
-            val = B[i, color[j]] - buffer_right_type[i]
-            buffer_right_type[j] = buffer_right_type[j] + val
+            if !is_star[k]
+                # the tree is not a star
+                val = B[i, color[j]] - buffer_right_type[i]
+                buffer_right_type[j] = buffer_right_type[j] + val
+            else
+                # the tree is a trivial or non-trivial star
+                val = B[i, color[j]]
+            end
             if in_triangle(i, j, uplo)
                 A[i, j] = val
             end
@@ -554,6 +563,7 @@ function decompress!(
         color,
         vertices_by_tree,
         reverse_bfs_orders,
+        is_star,
         diagonal_indices,
         diagonal_nzind,
         lower_triangle_offsets,
@@ -596,14 +606,23 @@ function decompress!(
 
     # Recover the off-diagonal coefficients of A
     for k in eachindex(vertices_by_tree, reverse_bfs_orders)
-        for vertex in vertices_by_tree[k]
-            buffer_right_type[vertex] = zero(R)
+        # We need the buffer only when the tree is not a star (trivial or non-trivial)
+        if !is_star[k]
+            for vertex in vertices_by_tree[k]
+                buffer_right_type[vertex] = zero(R)
+            end
         end
 
         for (i, j) in reverse_bfs_orders[k]
             counter += 1
-            val = B[i, color[j]] - buffer_right_type[i]
-            buffer_right_type[j] = buffer_right_type[j] + val
+            if !is_star[k]
+                # the tree is not a star
+                val = B[i, color[j]] - buffer_right_type[i]
+                buffer_right_type[j] = buffer_right_type[j] + val
+            else
+                # the tree is a trivial or non-trivial star
+                val = B[i, color[j]]
+            end
 
             #! format: off
             # A[i,j] is in the lower triangular part of A

src/result.jl

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