preorder and postorder can take a list of roots

kkowar · Sep 24, 2014 · 2a9d8ca · 2a9d8ca
1 parent b90f9bf
commit 2a9d8ca
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 11 deletions.
diff --git a/lib/alg/dfs.js b/lib/alg/dfs.js
@@ -3,7 +3,7 @@ var _ = require("lodash");
 module.exports = dfs;
 
 /*
- * dfs(graph, root, order, [traversalFunction], [successorFunction])
+ * dfs(graph, roots, order, [traversalFunction], [successorFunction])
  *
  * A helper function for doing dfs traversals in pre- or post-order. The root
  * node must be in the graph.
@@ -19,21 +19,27 @@ module.exports = dfs;
  * that should be traversed. By default it uses the graph's "successor"
  * function.
  */
-function dfs(g, v, order, fn, successors) {
+function dfs(g, vs, order, fn, successors) {
   if (!successors) {
     successors = g.successors.bind(g);
   }
 
-  if (!g.hasNode(v)) {
-    throw new Error("Graph does not have node: " + v);
-  }
-
   if (!fn) {
     fn = _.constant(true);
   }
 
+  if (!_.isArray(vs)) {
+    vs = [vs];
+  }
+
   var acc = [];
-  doDfs(v, order, fn, successors, {}, acc);
+  _.each(vs, function(v) {
+    if (!g.hasNode(v)) {
+      throw new Error("Graph does not have node: " + v);
+    }
+
+    doDfs(v, order, fn, successors, {}, acc);
+  });
   return acc;
 }
 

diff --git a/lib/alg/postorder.js b/lib/alg/postorder.js
@@ -2,6 +2,6 @@ var dfs = require("./dfs");
 
 module.exports = preorder;
 
-function preorder(g, v, fn, successors) {
-  return dfs(g, v, "post", fn, successors);
+function preorder(g, vs, fn, successors) {
+  return dfs(g, vs, "post", fn, successors);
 }
diff --git a/lib/alg/preorder.js b/lib/alg/preorder.js
@@ -2,6 +2,6 @@ var dfs = require("./dfs");
 
 module.exports = preorder;
 
-function preorder(g, v, fn, successors) {
-  return dfs(g, v, "pre", fn, successors);
+function preorder(g, vs, fn, successors) {
+  return dfs(g, vs, "pre", fn, successors);
 }
diff --git a/test/alg/postorder-test.js b/test/alg/postorder-test.js
@@ -62,6 +62,19 @@ describe("alg.postorder", function() {
         Math.min(nodes.indexOf("b"), nodes.indexOf("c")));
   });
 
+  it("works for an array of roots", function() {
+    var g = new Digraph();
+    g.setEdge("a", "b");
+    g.setEdge("c", "d");
+    g.setNode("e");
+    g.setNode("f");
+
+    var nodes = postorder(g, ["a", "c", "e"]);
+    expect(_.sortBy(nodes)).to.eql(["a", "b", "c", "d", "e"]);
+    expect(nodes.indexOf("b")).to.be.lt(nodes.indexOf("a"));
+    expect(nodes.indexOf("d")).to.be.lt(nodes.indexOf("c"));
+  });
+
   it("finishes early if the callback returns false", function() {
     var g = new Digraph();
     g.setPath(["a", "b", "c", "d"]);

diff --git a/test/alg/preorder-test.js b/test/alg/preorder-test.js
@@ -64,6 +64,19 @@ describe("alg.preorder", function() {
         nodes.indexOf("d"));
   });
 
+  it("works for an array of roots", function() {
+    var g = new Digraph();
+    g.setEdge("a", "b");
+    g.setEdge("c", "d");
+    g.setNode("e");
+    g.setNode("f");
+
+    var nodes = preorder(g, ["a", "c", "e"]);
+    expect(_.sortBy(nodes)).to.eql(["a", "b", "c", "d", "e"]);
+    expect(nodes.indexOf("b")).to.be.gt(nodes.indexOf("a"));
+    expect(nodes.indexOf("d")).to.be.gt(nodes.indexOf("c"));
+  });
+
   it("finishes early if the callback returns false", function() {
     var g = new Digraph();
     g.setPath(["a", "b", "c", "d"]);