Merge pull request #31 from fmjrey/edge-traverse

Edge traverse

src/loom/alg_generic.clj

@@ -62,6 +62,30 @@
     (step [start]
           (conj seen start))))
 
+(defn pre-edge-traverse
+  "Traverses a graph depth-first preorder from start, successors being
+  a function that returns direct successors for the node. Returns a
+  lazy seq of edges, each edge being a vector [source-node dest-node].
+  Note that for undirected graphs each edge will be returned twice,
+  once for each direction."
+  [successors start & {:keys [seen] :or {seen #{}}}]
+  (letfn [(step [successors start nbrs stack nbrstack seen]
+    (if-let [nbr (first nbrs)]
+      (cons
+        [start nbr]
+        (lazy-seq
+          (let [seen (conj seen start)]
+            (if (seen nbr)
+              (step successors start (next nbrs) stack nbrstack seen)
+              (step successors nbr (successors nbr)
+                    (conj stack start) (conj nbrstack (next nbrs))
+                    seen)))))
+      (when-let [parent (peek stack)]
+        (recur successors parent (peek nbrstack)
+               (pop stack) (pop nbrstack) (conj seen start)))))] 
+    (when-not (seen start)
+      (step successors start (successors start) [] [] (conj seen start)))))
+
 ;; TODO: graph-seq, analog of tree-seq
 
 (defn pre-span
@@ -100,6 +124,37 @@
           (recur seen (conj result v) (pop stack))
           (recur seen result (conj stack (first nbrs))))))))
 
+(defn post-edge-traverse
+  "Traverses a graph depth-first postorder from start, successors being
+  a function that returns direct successors for the node. Returns a
+  seq of edges, each edge being a vector [source-node dest-node].
+  Note that for undirected graphs each edge will be returned twice,
+  once for each direction."
+  [successors start & {:keys [seen return-seen] :or {seen #{}}}]
+  (if (seen start)
+    (if return-seen
+      [nil seen]
+      nil)
+    (loop [start start
+          nbrs (successors start)
+          stack []
+          nbrstack []
+          seen seen
+          edges ()]
+      (let [seen (conj seen start)]
+        (if-let [nbr (first nbrs)]
+          (if (seen nbr)
+            (recur start (next nbrs) stack nbrstack seen (conj edges [start nbr]))
+            (recur nbr (successors nbr)
+                   (conj stack start) (conj nbrstack (next nbrs))
+                   seen (conj edges [start nbr])))
+          (if-let [parent (peek stack)]
+            (recur parent (peek nbrstack)
+                   (pop stack) (pop nbrstack) seen edges)
+            (if return-seen
+              [edges seen]
+              edges)))))))
+
 (defn topsort-component
   "Topological sort of a component of a (presumably) directed graph.
   Returns nil if the graph contains any cycles. See loom.alg/topsort

test/loom/test/alg_generic.clj

@@ -1,5 +1,6 @@
 (ns loom.test.alg-generic
   (:require [loom.alg-generic :as lag]
+            [loom.graph :as g]
             [clojure.set :as set]
             [clojure.test :refer [deftest are]]
             [clojure.test.check :as tc]
@@ -114,6 +115,28 @@
    :e [:c :d :f]
    :f []})
 
+(def g4 ; like g3 with some loops
+  {:a [:b]
+   :b [:a :c :d]
+   :c [:b :c :e]
+   :d [:b :c :e]
+   :e [:c :d :f]
+   :f [:f]})
+
+(def g5 ; like g1 but as an undirected graph
+  {:a [:b :c]
+   :b [:d :a]
+   :c [:a :d]
+   :d [:c :b]})
+
+(def g6 ; unconnected with some loops
+  {:a [:a]
+   :b [:a :c]
+   :c [:b :c]
+   :d [:e]
+   :e [:d :f]
+   :f [:f]})
+
 (deftest tracing-paths
   (are [g n p] (= (sort (lag/trace-paths g n)) p)
        {:a nil} :a
@@ -135,3 +158,56 @@
 
        g3 :a :e
        [[:a :b :c :e] [:a :b :d :e]]))
+
+(deftest edge-traverse
+  ; works with nodes without outgoing edges or just a loop to iself
+  (are [g start expected] (let [pre (lag/pre-edge-traverse g start)
+                                post (lag/post-edge-traverse g start)]
+                            (= expected pre (seq (reverse post))))
+       g1 :d nil
+
+       g4 :f '([:f :f]))
+  ; covers the whole graph when it's totally connected from start
+  (are [g start expected] (let [pre (lag/pre-edge-traverse g start)
+                                post (lag/post-edge-traverse g start)
+                                dg (g/digraph g)
+                                edges (g/edges dg)]
+                            (and
+                              (= expected pre (seq (reverse post)))
+                              (= (count edges) (count post))
+                              (= (set edges) (set post))))
+       g1 :a '([:a :b] [:b :d] [:a :c] [:c :d])
+
+       g4 :a '([:a :b] [:b :a] [:b :c] [:c :b] [:c :c] [:c :e] [:e :c]
+               [:e :d] [:d :b] [:d :c] [:d :e] [:e :f] [:f :f] [:b :d])
+
+       g4 :c '([:c :b] [:b :a] [:a :b] [:b :c] [:b :d] [:d :b] [:d :c]
+               [:d :e] [:e :c] [:e :d] [:e :f] [:f :f] [:c :c] [:c :e])
+
+       g5 :a '([:a :b] [:b :d] [:d :c] [:c :a]
+               [:c :d] [:d :b] [:b :a] [:a :c]))
+  ; post traversal returning seen nodes allows complete graph coverage
+  ; without duplicates when iterating on all nodes of the graph
+  (are [g] (let [dg (g/digraph g)
+                 edges (g/edges dg)
+                 loop-post-traverse
+                     (loop [nodes (reverse (g/nodes dg))
+                            ; reverse makes this more interesting as graphs
+                            ; are often specified in the forward direction
+                            seen #{}
+                            acc ()]
+                       (if-let [node (first nodes)]
+                         (let [[edges seen]
+                                   (lag/post-edge-traverse
+                                     g
+                                     node
+                                     :seen seen
+                                     :return-seen true)]
+                           (recur (next nodes)
+                                  seen
+                                  (concat acc edges)))
+                         acc))]
+             (and
+               (= (count edges) (count loop-post-traverse))
+               (= (set edges) (set loop-post-traverse))))
+       g1 g2 g3 g4 g5 g6))