Faster all_neighbors for SimpleDiGraphs (#82)

Author: etiennedeg

src/SimpleGraphs/SimpleGraphs.jl

@@ -11,8 +11,9 @@ import Base:
 import Graphs:
     _NI, AbstractGraph, AbstractEdge, AbstractEdgeIter,
     src, dst, edgetype, nv, ne, vertices, edges, is_directed,
-    has_vertex, has_edge, inneighbors, outneighbors, deepcopy_adjlist,
-    indegree, outdegree, degree, has_self_loops, num_self_loops, insorted
+    has_vertex, has_edge, inneighbors, outneighbors, all_neighbors,
+    deepcopy_adjlist, indegree, outdegree, degree, has_self_loops,
+    num_self_loops, insorted
 
 using Random: GLOBAL_RNG, AbstractRNG
 

src/SimpleGraphs/simpledigraph.jl

@@ -566,3 +566,37 @@ function rem_vertices!(g::SimpleDiGraph{T},
 
     return reverse_vmap
 end
+
+function all_neighbors(g::SimpleDiGraph{T}, u::Integer) where T
+    i, j = 1, 1
+    in_nbrs, out_nbrs = inneighbors(g, u), outneighbors(g, u)
+    in_len, out_len = length(in_nbrs), length(out_nbrs)
+    union_nbrs = Vector{T}(undef, in_len + out_len)
+    indx = 1
+    @inbounds while i <= in_len && j <= out_len
+        if in_nbrs[i] < out_nbrs[j]
+            union_nbrs[indx] = in_nbrs[i]
+            i += 1
+        elseif in_nbrs[i] > out_nbrs[j]
+            union_nbrs[indx] = out_nbrs[j]
+            j += 1
+        else 
+            union_nbrs[indx] = out_nbrs[j]
+            i += 1
+            j += 1
+        end
+        indx += 1
+    end
+    @inbounds while i <= in_len
+        union_nbrs[indx] = in_nbrs[i]
+        i += 1
+        indx += 1
+    end
+    @inbounds while j <= out_len
+        union_nbrs[indx] = out_nbrs[j]
+        j += 1
+        indx += 1
+    end
+    resize!(union_nbrs, indx-1)
+    return union_nbrs
+end

test/simplegraphs/simplegraphs.jl

@@ -129,6 +129,7 @@ import Random
 
         @test @inferred(inneighbors(g, 2)) == [1]
         @test @inferred(outneighbors(g, 2)) == @inferred(neighbors(g, 2)) == [3]
+        @test @inferred Set(all_neighbors(g, 2)) == Set(union(inneighbors(g, 2), outneighbors(g, 2)))
         @test @inferred(add_vertex!(gc))   # out of order, but we want it for issubset
         @test @inferred(g ⊆ gc)
         @test @inferred(has_vertex(gc, 5))