relax colormap type in traversal

ropenta · Apr 12, 2016 · d522ed1 · d522ed1
1 parent 4a09f3b
commit d522ed1
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 19 deletions.
diff --git a/src/traversals/bfs.jl b/src/traversals/bfs.jl
@@ -13,21 +13,19 @@ type BreadthFirst <: SimpleGraphVisitAlgorithm
 end
 
 function breadth_first_visit_impl!(
-    graph::SimpleGraph,   # the graph
-    queue::Vector{Int},                  # an (initialized) queue that stores the active vertices
-    colormap::Vector{Int},          # an (initialized) color-map to indicate status of vertices
-    visitor::SimpleGraphVisitor)  # the visitor
+    graph::SimpleGraph,                 # the graph
+    queue::Vector{Int},                 # an (initialized) queue that stores the active vertices
+    colormap,                           # an (initialized) color-map to indicate status of vertices (0=unseen, 1=seen, 2=closed)
+    visitor::SimpleGraphVisitor)        # the visitor
 
     while !isempty(queue)
         u = shift!(queue)
         open_vertex!(visitor, u)
 
         for v in out_neighbors(graph, u)
-            v_color::Int = colormap[v]
+            v_color = get(colormap, v, 0)
             # TODO: Incorporate edge colors to BFS
-            if !(examine_neighbor!(visitor, u, v, v_color, -1))
-                return
-            end
+            examine_neighbor!(visitor, u, v, v_color, -1) || return
 
             if v_color == 0
                 colormap[v] = 1

diff --git a/src/traversals/dfs.jl b/src/traversals/dfs.jl
@@ -16,7 +16,7 @@ end
 function depth_first_visit_impl!(
     graph::SimpleGraph,      # the graph
     stack,                          # an (initialized) stack of vertex
-    vertexcolormap::Vector{Int},    # an (initialized) color-map to indicate status of vertices
+    vertexcolormap,    # an (initialized) color-map to indicate status of vertices
     edgecolormap::Dict{Edge,Int},      # an (initialized) color-map to indicate status of edges
     visitor::SimpleGraphVisitor)  # the visitor
 
@@ -27,7 +27,7 @@ function depth_first_visit_impl!(
 
         while !done(udsts, tstate) && !found_new_vertex
             v, tstate = next(udsts, tstate)
-            v_color = vertexcolormap[v]
+            v_color = get(vertexcolormap, v, 0)
             v_edge = Edge(u,v)
             e_color = haskey(edgecolormap, v_edge)?
                 edgecolormap[v_edge] : edgecolormap[reverse(v_edge)]

diff --git a/test/traversals/bfs.jl b/test/traversals/bfs.jl
@@ -18,15 +18,8 @@ add_edge!(g,3,4)
 
 
 import LightGraphs: TreeBFSVisitorVector, bfs_tree!, TreeBFSVisitor, tree
-g = smallgraph(:house)
-n = nv(g)
-visitor = TreeBFSVisitorVector(n)
-@test length(visitor.tree) == n
-parents = visitor.tree
-bfs_tree!(visitor, g, 1)
-maxdepth = n
 
-function istree(parents::Vector{Int})
+function istree(parents::Vector{Int}, maxdepth)
     flag = true
     for i in 1:n
         s = i
@@ -42,7 +35,32 @@ function istree(parents::Vector{Int})
     return flag
 end
 
-@test istree(parents) == true
+g = smallgraph(:house)
+n = nv(g)
+visitor = TreeBFSVisitorVector(n)
+@test length(visitor.tree) == n
+parents = visitor.tree
+bfs_tree!(visitor, g, 1)
+
+@test istree(parents, n) == true
+tvis = TreeBFSVisitor(visitor)
+@test nv(tvis.tree) == nv(g)
+@test typeof(tvis.tree) <: DiGraph
+t = tree(parents)
+@test typeof(t) <: DiGraph
+@test typeof(tvis.tree) <: DiGraph
+@test t == tvis.tree
+@test ne(t) < nv(t)
+
+# test Dict{Int,Int}() colormap
+g = smallgraph(:house)
+n = nv(g)
+visitor = TreeBFSVisitorVector(n)
+@test length(visitor.tree) == n
+parents = visitor.tree
+bfs_tree!(visitor, g, 1, colormap = Dict{Int,Int}())
+
+@test istree(parents, n) == true
 tvis = TreeBFSVisitor(visitor)
 @test nv(tvis.tree) == nv(g)
 @test typeof(tvis.tree) <: DiGraph