More optimization for the decompression of the acyclic coloring (#73)

* More optimization for the decompression of the acyclic coloring

* Optimal version!

* Reduce dynamic allocations

* Avoid modifying `TreeSet` a posteriori

* Typo

* Renaming

* Typo

---------

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

Author: amontoison

src/coloring.jl

@@ -328,6 +328,12 @@ function acyclic_coloring(g::Graph, order::AbstractOrder)
             end
         end
     end
+
+    # compress forest
+    for edge in forest.revmap
+        find_root!(forest, edge)
+    end
+
     return color, TreeSet(forest)
 end
 
@@ -343,8 +349,8 @@ function _prevent_cycle!(
     forest::DisjointSets{<:Tuple{Int,Int}},
 )
     wx = _sort(w, x)
-    root = find_root!(forest, wx)  # edge wx belongs to the 2-colored tree represented by edge "root"
-    id = forest.intmap[root] # ID of the representative edge "root" of a two-colored tree.
+    root = find_root!(forest, wx)  # edge wx belongs to the 2-colored tree T represented by edge "root"
+    id = forest.intmap[root] # ID of the representative edge "root" of a two-colored tree T.
     (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)

src/decompression.jl

@@ -288,14 +288,12 @@ function decompress_aux!(
     A .= zero(R)
     S = get_matrix(result)
     color = column_colors(result)
-    @compat (; degrees, dfs_orders, stored_values) = result
+    @compat (; vertices_by_tree, reverse_bfs_orders, buffer) = result
 
-    # stored_values holds the sum of edge values for subtrees in a tree.
-    # For each vertex i, stored_values[i] is the sum of edge values in the subtree rooted at i.
-    stored_values_right_type = if R == Float64
-        stored_values
+    if eltype(buffer) == R
+        buffer_right_type = buffer
     else
-        similar(stored_values, R)
+        buffer_right_type = similar(buffer, R)
     end
 
     # Recover the diagonal coefficients of A
@@ -306,16 +304,14 @@ function decompress_aux!(
     end
 
     # Recover the off-diagonal coefficients of A
-    for k in eachindex(degrees, dfs_orders)
-        vertices = keys(degrees[k])
-        for vertex in vertices
-            stored_values_right_type[vertex] = zero(R)
+    for k in eachindex(vertices_by_tree, reverse_bfs_orders)
+        for vertex in vertices_by_tree[k]
+            buffer_right_type[vertex] = zero(R)
         end
 
-        tree = dfs_orders[k]
-        for (i, j) in tree
-            val = B[i, color[j]] - stored_values_right_type[i]
-            stored_values_right_type[j] = stored_values_right_type[j] + val
+        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
             A[i, j] = val
             A[j, i] = val
         end

src/result.jl

@@ -198,15 +198,15 @@ struct TreeSetColoringResult{M,R} <:
     S::M
     color::Vector{Int}
     group::Vector{Vector{Int}}
-    tree_set::TreeSet
-    degrees::Vector{Dict{Int,Int}}
-    dfs_orders::Vector{Vector{Tuple{Int,Int}}}
-    stored_values::Vector{R}
+    vertices_by_tree::Vector{Vector{Int}}
+    reverse_bfs_orders::Vector{Vector{Tuple{Int,Int}}}
+    buffer::Vector{R}
 end
 
 function TreeSetColoringResult(
     S::SparseMatrixCSC, color::Vector{Int}, tree_set::TreeSet, decompression_eltype::Type{R}
 ) where {R}
+    nvertices = length(color)
     group = group_by_color(color)
 
     # forest is a structure DisjointSets from DataStructures.jl
@@ -216,66 +216,101 @@ function TreeSetColoringResult(
     forest = tree_set.forest
     ntrees = forest.internal.ngroups
 
-    # vector of trees where each tree contains the indices of its edges
-    trees = [Int[] for i in 1:ntrees]
-
     # dictionary that maps a tree's root to the index of the tree
     roots = Dict{Int,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:ntrees]
+
+    # counter of the number of roots found
     k = 0
     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]
+
+        # 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]
-        push!(trees[index_tree], forest.intmap[edge])
-    end
 
-    # vector of dictionaries where each dictionary stores the degree of each vertex in a tree
-    degrees = [Dict{Int,Int}() for k in 1:ntrees]
-    for k in 1:ntrees
-        tree = trees[k]
-        degree = degrees[k]
-        for edge_index in tree
-            i, j = forest.revmap[edge_index]
-            !haskey(degree, i) && (degree[i] = 0)
-            !haskey(degree, j) && (degree[j] = 0)
-            degree[i] += 1
-            degree[j] += 1
+        # 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
 
-    # depth-first search (DFS) traversal order for each tree in the forest
-    dfs_orders = [Vector{Tuple{Int,Int}}() for k in 1:ntrees]
+    # 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]
+
     for k in 1:ntrees
         tree = trees[k]
-        degree = degrees[k]
-        while sum(values(degree)) != 0
-            for (t, edge_index) in enumerate(tree)
-                if edge_index != 0
-                    i, j = forest.revmap[edge_index]
-                    if (degree[i] == 1) || (degree[j] == 1)  # leaf vertex
-                        if degree[i] > degree[j]  # vertex i is the parent of vertex j
-                            i, j = j, i  # ensure that i always denotes a leaf vertex
-                        end
-                        degree[i] -= 1  # decrease the degree of vertex i
-                        degree[j] -= 1  # decrease the degree of vertex j
-                        tree[t] = 0  # remove the edge (i,j)
-                        push!(dfs_orders[k], (i, j))
+
+        # queue to store the leaves
+        queue = Int[]
+
+        # compute the degree of each vertex in the tree
+        for (vertex, neighbors) in trees[k]
+            degree = length(neighbors)
+            degrees[vertex] = degree
+
+            # the vertex is a leaf
+            if degree == 1
+                push!(queue, vertex)
+            end
+        end
+
+        # continue until all leaves are treated
+        while !isempty(queue)
+            leaf = pop!(queue)
+
+            # Convenient way to specify that the vertex is removed
+            degrees[leaf] = 0
+
+            for neighbor in tree[leaf]
+                if degrees[neighbor] != 0
+                    # (leaf, neighbor) represents the next edge to visit during decompression
+                    push!(reverse_bfs_orders[k], (leaf, neighbor))
+
+                    # reduce the degree of all neighbors
+                    degrees[neighbor] -= 1
+
+                    # check if the neighbor is now a leaf
+                    if degrees[neighbor] == 1
+                        push!(queue, neighbor)
                     end
                 end
             end
         end
     end
 
-    n = checksquare(S)
-    stored_values = Vector{R}(undef, n)
+    # buffer holds the sum of edge values for subtrees in a tree.
+    # For each vertex i, buffer[i] is the sum of edge values in the subtree rooted at i.
+    buffer = Vector{R}(undef, nvertices)
 
     return TreeSetColoringResult(
-        S, color, group, tree_set, degrees, dfs_orders, stored_values
+        S, color, group, vertices_by_tree, reverse_bfs_orders, buffer
     )
 end