speed up full-graph traversals & spans, connected-components

src/loom/alg.clj

@@ -4,12 +4,11 @@ can use these functions."
       :author "Justin Kramer"}
   loom.alg
   (:require [loom.alg-generic :as gen])
-  (:use [clojure.set :only [union]]
-        [loom.graph
+  (:use [loom.graph
          :only [add-edges nodes edges neighbors weight incoming degree
                 in-degree weighted? directed? graph]
          :rename {neighbors nb weight wt}]
-        [loom.alg-generic :only [trace-path]]))
+        [loom.alg-generic :only [trace-path preds->span]]))
 
 ;;;
 ;;; Convenience wrappers for loom.alg-generic functions
@@ -44,8 +43,8 @@ can use these functions."
        (fn [[seen span] n]
          (if (seen n)
            [seen span]
-           (let [[cspan cseen] (gen/pre-span (nb g) n :seen seen)]
-             [(union seen cseen) (merge span {n []} cspan)])))
+           (let [[cspan seen] (gen/pre-span (nb g) n :seen seen :return-seen true)]
+             [seen (merge span {n []} cspan)])))
        [#{} {}]
        (nodes g))))
   ([g start]
@@ -81,7 +80,18 @@ can use these functions."
   Otherwise, returns a lazy seq of the nodes. When option :when is provided,
   filters neighbors with (f neighbor predecessor depth)."
   ([g]
-     (traverse-all (nodes g) (partial gen/bf-traverse (nb g))))
+     (first
+      (reduce
+       (fn [[cc predmap] n]
+         (if (contains? predmap n)
+           [cc predmap]
+           (reduce
+            (fn [[cc _] [n pm _]]
+              [(conj cc n) pm])
+            [cc predmap]
+            (gen/bf-traverse (nb g) n :f vector :seen predmap))))
+       [[] {}]
+       (nodes g))))
   ([g start]
      (gen/bf-traverse (nb g) start))
   ([g start & opts]
@@ -90,16 +100,15 @@ can use these functions."
 (defn bf-span
   "Return a breadth-first spanning tree of the form {node [successors]}"
   ([g]
-     (second
+     (preds->span
       (reduce
-       (fn [[seen span] n]
-         (if (seen n)
-           [seen span]
-           (let [cspan (gen/bf-span (nb g) n :seen seen)]
-             ;; FIXME: very inefficient
-             [(into seen (concat (keys cspan) (apply concat (vals cspan))))
-              (merge span {n []} cspan)])))
-       [#{} {}]
+       (fn [predmap n]
+         (if (contains? predmap n)
+           predmap
+           (last (gen/bf-traverse (nb g) n
+                                  :f (fn [_ pm _] pm)
+                                  :seen predmap))))
+       {}
        (nodes g))))
   ([g start]
      (gen/bf-span (nb g) start)))
@@ -184,15 +193,18 @@ can use these functions."
 (defn connected-components
   "Return the connected components of undirected graph g as a vector of vectors"
   [g]
-  ;; TODO: leverage predmap from bf-traverse (keys = seen)
   (first
    (reduce
-    (fn [[cc seen] n]
-      (if (seen n)
-        [cc seen]
-        (let [c (vec (gen/bf-traverse (nb g) n :seen seen))]
-          [(conj cc c) (into seen c)])))
-    [[] #{}]
+    (fn [[cc predmap] n]
+      (if (contains? predmap n)
+        [cc predmap]
+        (let [[c pm] (reduce
+                      (fn [[c _] [n pm _]]
+                        [(conj c n) pm])
+                      [[] nil]
+                      (gen/bf-traverse (nb g) n :f vector :seen predmap))]
+          [(conj cc c) pm])))
+    [[] {}]
     (nodes g))))
 
 ;; TODO: weak & strong cc

src/loom/alg_generic.clj

@@ -30,7 +30,7 @@
 (defn pre-traverse
   "Traverses a graph depth-first preorder from start, neighbors being a
   function that returns adjacent nodes. Returns a lazy seq of nodes."
-  [neighbors start & {:keys [seen] :or {seen #{start}}}]
+  [neighbors start & {:keys [seen] :or {seen #{}}}]
   (letfn [(step [stack seen]
             (when-let [node (peek stack)]
               (cons
@@ -40,24 +40,24 @@
                   (step (into (pop stack) nbrs)
                         (into seen nbrs)))))))]
     (step [start]
-          seen)))
+          (conj seen start))))
 
 (defn pre-span
   "Return a depth-first spanning tree of the form {node [successors]}"
-  [neighbors start & {:keys [seen]}]
-  (let [seen-start seen]
-    (loop [seen (or seen-start #{})
-           preds {start nil}
-           stack [start]]
-      (if (empty? stack)
-        (if seen-start
-          [(preds->span preds) seen]
-          (preds->span preds))
-        (let [v (peek stack)
-              seen (conj seen v)]
-          (if-let [u (first (remove seen (neighbors v)))]
-            (recur seen (assoc preds u v) (conj stack u))
-            (recur seen preds (pop stack))))))))
+  [neighbors start & {:keys [seen return-seen] :or {seen #{}}}]
+  (loop [seen seen
+         preds {start nil}
+         stack [start]]
+    (if (empty? stack)
+      ;; TODO: this is awkward, devise something better
+      (if return-seen
+        [(preds->span preds) seen]
+        (preds->span preds))
+      (let [v (peek stack)
+            seen (conj seen v)]
+        (if-let [u (first (remove seen (neighbors v)))]
+          (recur seen (assoc preds u v) (conj stack u))
+          (recur seen preds (pop stack)))))))
 
 (defn post-traverse
   "Traverses a graph depth-first postorder from start, neighbors being a
@@ -123,17 +123,18 @@
                     (step (into (pop queue) (for [nbr nbrs] [nbr (inc depth)]))
                           (reduce #(assoc %1 %2 node) preds nbrs)))))))]
       (step (conj clojure.lang.PersistentQueue/EMPTY [start 0])
-            (into {start nil} (if (map? seen)
-                                seen
-                                (for [s seen] [s nil])))))))
+            (if (map? seen)
+              (assoc seen start nil)
+              (into {start nil} (for [s seen] [s nil])))))))
 
 (defn bf-span
   "Return a breadth-first spanning tree of the form {node [successors]}"
   [neighbors start & {:keys [seen]}]
   (preds->span
-   (bf-traverse neighbors start
-                :f (fn [n pm _] [n (pm n)])
-                :seen seen)))
+   (last
+    (bf-traverse neighbors start
+                 :f (fn [_ pm _] pm)
+                 :seen seen))))
 
 (defn bf-path
   "Return a path from start to end with the fewest hops (i.e. irrespective