Add unit tests for Forest

Author: amontoison

src/coloring.jl

@@ -432,15 +432,15 @@ 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.ntrees: the number of trees in the forest
-    ntrees = forest.ntrees
+    # - forest.num_trees: the number of trees in the forest
+    nt = forest.num_trees
 
     # dictionary that maps a tree's root to the index of the tree
     roots = Dict{Int,Int}()
-    sizehint!(roots, ntrees)
+    sizehint!(roots, nt)
 
     # 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]
+    trees = [Dict{Int,Vector{Int}}() for i in 1:nt]
 
     # counter of the number of roots found
     k = 0
@@ -476,11 +476,11 @@ function TreeSet(forest::Forest{Int}, nvertices::Int)
     degrees = Vector{Int}(undef, nvertices)
 
     # reverse breadth first (BFS) traversal order for each tree in the forest
-    reverse_bfs_orders = [Tuple{Int,Int}[] for i in 1:ntrees]
+    reverse_bfs_orders = [Tuple{Int,Int}[] for i in 1:nt]
 
     # nvmax is the number of vertices of the biggest tree in the forest
     nvmax = 0
-    for k in 1:ntrees
+    for k in 1:nt
         nb_vertices_tree = length(trees[k])
         nvmax = max(nvmax, nb_vertices_tree)
     end
@@ -490,9 +490,9 @@ function TreeSet(forest::Forest{Int}, nvertices::Int)
 
     # 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)
+    is_star = Vector{Bool}(undef, nt)
 
-    for k in 1:ntrees
+    for k in 1:nt
         tree = trees[k]
 
         # Boolean indicating whether the current tree is a star (a single central vertex connected to all others)

src/forest.jl

@@ -10,32 +10,32 @@ Structure that provides fast union-find operations for constructing a forest dur
 $TYPEDFIELDS
 """
 mutable struct Forest{T<:Integer}
-    "current number of edges added to the forest"
-    counter::T
+    "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}
-    "current number of distinct trees in the forest"
-    ntrees::T
 end
 
 function Forest{T}(n::Integer) where {T<:Integer}
-    counter = zero(T)
+    num_edges = zero(T)
+    num_trees = zero(T)
     intmap = Dict{Tuple{T,T},T}()
     sizehint!(intmap, n)
     parents = collect(Base.OneTo(T(n)))
     ranks = zeros(T, T(n))
-    ntrees = zero(T)
-    return Forest{T}(counter, intmap, parents, ranks, ntrees)
+    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.counter += 1
-    forest.intmap[edge] = forest.counter
-    forest.ntrees += one(T)
+    forest.num_edges += 1
+    forest.intmap[edge] = forest.num_edges
+    forest.num_trees += one(T)
     return forest
 end
 
@@ -63,6 +63,6 @@ function root_union!(forest::Forest{T}, index_edge1::T, index_edge2::T) where {T
         rks[index_edge1] += one(T)
     end
     parents[index_edge2] = index_edge1
-    forest.ntrees -= one(T)
+    forest.num_trees -= one(T)
     return nothing
 end

test/forest.jl

@@ -0,0 +1,73 @@
+@testset "Constructor Forest" begin
+    forest = Forest{Int}(5)
+
+    @test forest.num_edges == 0
+    @test forest.num_trees == 0
+    @test length(forest.intmap) == 0
+    @test length(forest.parents) == 5
+    @test all(forest.parents .== 1:5)
+    @test all(forest.ranks .== 0)
+end
+
+@testset "Push edge" begin
+    forest = Forest{Int}(5)
+
+    push!(forest, (1, 2))
+    @test forest.num_edges == 1
+    @test forest.num_trees == 1
+    @test haskey(forest.intmap, (1, 2))
+    @test forest.intmap[(1, 2)] == 1
+    @test forest.num_trees == 1
+
+    push!(forest, (3, 4))
+    @test forest.num_edges == 2
+    @test forest.num_trees == 2
+    @test haskey(forest.intmap, (3, 4))
+    @test forest.intmap[(3, 4)] == 2
+    @test forest.num_trees == 2
+end
+
+@testset "Find root" begin
+    forest = Forest{Int}(5)
+    push!(forest, (1, 2))
+    push!(forest, (3, 4))
+
+    @test find_root!(forest, (1, 2)) == 1
+    @test find_root!(forest, (3, 4)) == 2
+end
+
+@testset "Root union" begin
+    forest = Forest{Int}(5)
+    push!(forest, (1, 2))
+    push!(forest, (4, 5))
+    push!(forest, (2, 4))
+    @test forest.num_trees = 3
+
+    root1 = find_root!(forest, (1, 2))
+    root3 = find_root!(forest, (2, 4))
+    @test root1 != root3
+
+    root_union!(forest, root1, root3)
+    @test find_root!(forest, (2,4)) == 1
+    @test forest.parents[1] == 1
+    @test forest.parents[3] == 1
+    @test forest.ranks[1] == 1
+    @test forest.ranks[3] == 0
+    @test forest.num_trees = 2
+
+    root1 = find_root!(forest, (1, 2))
+    root2 = find_root!(forest, (4, 5))
+    @test root1 != root2
+    root_union!(forest, root1, root2)
+    @test find_root!(forest, (4, 5)) == 1
+    @test forest.parents[1] == 1
+    @test forest.parents[2] == 1
+    @test forest.ranks[1] == 1
+    @test forest.ranks[2] == 0
+    @test forest.num_trees == 1
+
+    push!(forest, (1, 4))
+    @test forest.num_trees == 2
+    @test forest.intmap[(1, 4)] == 4
+    @test forest.parents[4] == 4
+end

test/runtests.jl

@@ -33,6 +33,9 @@ include("utils.jl")
         @testset "Graph" begin
             include("graph.jl")
         end
+        @testset "Forest" begin
+            include("forest.jl")
+        end
         @testset "Order" begin
             include("order.jl")
         end