search: avoid duplicated allocations for nodes

Also reduce code duplication.

search/dijkstra.go

@@ -6,10 +6,8 @@ package search
 
 import (
 	"container/heap"
-	"math"
 
 	"github.com/gonum/graph"
-	"github.com/gonum/matrix/mat64"
 )
 
 // DijkstraFrom returns a shortest-path tree for a shortest path from u to all nodes in
@@ -78,6 +76,16 @@ func DijkstraFrom(u graph.Node, g graph.Graph, weight graph.CostFunc) Shortest {
 //
 // The time complexity of DijkstrAllPaths is O(|V|.|E|+|V|^2.log|V|).
 func DijkstraAllPaths(g graph.Graph, weight graph.CostFunc) (paths AllShortest) {
+	paths = newAllShortest(g.NodeList(), false)
+	dijkstraAllPaths(g, weight, paths)
+	return paths
+}
+
+// dijkstraAllPaths is the all-paths implementation of Dijkstra. It is shared
+// between DijkstraAllPaths and JohnsonAllPaths to avoid repeated allocation
+// of the nodes slice and the indexOf map. It returns nothing, but stores the
+// result of the work in the paths parameter which is a reference type.
+func dijkstraAllPaths(g graph.Graph, weight graph.CostFunc, paths AllShortest) {
 	var (
 		from   = g.Neighbors
 		edgeTo func(graph.Node, graph.Node) graph.Edge
@@ -97,28 +105,8 @@ func DijkstraAllPaths(g graph.Graph, weight graph.CostFunc) (paths AllShortest)
 		}
 	}
 
-	nodes := g.NodeList()
-
-	indexOf := make(map[int]int, len(nodes))
-	for i, n := range nodes {
-		indexOf[n.ID()] = i
-	}
-
-	dist := make([]float64, len(nodes)*len(nodes))
-	for i := range dist {
-		dist[i] = math.Inf(1)
-	}
-	paths = AllShortest{
-		nodes:   nodes,
-		indexOf: indexOf,
-
-		dist:    mat64.NewDense(len(nodes), len(nodes), dist),
-		next:    make([][]int, len(nodes)*len(nodes)),
-		forward: false,
-	}
-
 	var Q priorityQueue
-	for i, u := range nodes {
+	for i, u := range paths.nodes {
 		// Dijkstra's algorithm here is implemented essentially as
 		// described in Function B.2 in figure 6 of UTCS Technical
 		// Report TR-07-54 with the addition of handling multiple
@@ -150,8 +138,6 @@ func DijkstraAllPaths(g graph.Graph, weight graph.CostFunc) (paths AllShortest)
 			}
 		}
 	}
-
-	return paths
 }
 
 type distanceNode struct {

search/floydwarshall.go

@@ -4,12 +4,7 @@
 
 package search
 
-import (
-	"math"
-
-	"github.com/gonum/graph"
-	"github.com/gonum/matrix/mat64"
-)
+import "github.com/gonum/graph"
 
 // FloydWarshall returns a shortest-path tree for the graph g or false indicating
 // that a negative cycle exists in the graph. If weight is nil and the graph does not
@@ -37,28 +32,11 @@ func FloydWarshall(g graph.Graph, weight graph.CostFunc) (paths AllShortest, ok
 	}
 
 	nodes := g.NodeList()
-
-	indexOf := make(map[int]int, len(nodes))
-	for i, n := range nodes {
-		indexOf[n.ID()] = i
-	}
-
-	dist := make([]float64, len(nodes)*len(nodes))
-	for i := range dist {
-		dist[i] = math.Inf(1)
-	}
-	paths = AllShortest{
-		nodes:   nodes,
-		indexOf: indexOf,
-
-		dist:    mat64.NewDense(len(nodes), len(nodes), dist),
-		next:    make([][]int, len(nodes)*len(nodes)),
-		forward: true,
-	}
+	paths = newAllShortest(nodes, true)
 	for i, u := range nodes {
 		paths.dist.Set(i, i, 0)
 		for _, v := range from(u) {
-			j := indexOf[v.ID()]
+			j := paths.indexOf[v.ID()]
 			paths.set(i, j, weight(edgeTo(u, v)), j)
 		}
 	}

search/johnson_apsp.go

@@ -39,17 +39,14 @@ func JohnsonAllPaths(g graph.Graph, weight graph.CostFunc) (paths AllShortest, o
 		}
 	}
 
-	nodes := g.NodeList()
-	indexOf := make(map[int]int, len(nodes))
-	for i, n := range nodes {
-		indexOf[n.ID()] = i
-	}
+	paths = newAllShortest(g.NodeList(), false)
+
 	sign := -1
 	for {
 		// Choose a random node ID until we find
 		// one that is not in g.
 		jg.q = sign * rand.Int()
-		if _, exists := indexOf[jg.q]; !exists {
+		if _, exists := paths.indexOf[jg.q]; !exists {
 			break
 		}
 		sign *= -1
@@ -62,7 +59,7 @@ func JohnsonAllPaths(g graph.Graph, weight graph.CostFunc) (paths AllShortest, o
 	}
 
 	jg.bellmanFord = false
-	paths = DijkstraAllPaths(jg, nil)
+	dijkstraAllPaths(jg, nil, paths)
 
 	for i, u := range paths.nodes {
 		hu := jg.adjustBy.WeightTo(u)

search/shortest.go

@@ -145,6 +145,25 @@ type AllShortest struct {
 	forward bool
 }
 
+func newAllShortest(nodes []graph.Node, forward bool) AllShortest {
+	indexOf := make(map[int]int, len(nodes))
+	for i, n := range nodes {
+		indexOf[n.ID()] = i
+	}
+	dist := make([]float64, len(nodes)*len(nodes))
+	for i := range dist {
+		dist[i] = math.Inf(1)
+	}
+	return AllShortest{
+		nodes:   nodes,
+		indexOf: indexOf,
+
+		dist:    mat64.NewDense(len(nodes), len(nodes), dist),
+		next:    make([][]int, len(nodes)*len(nodes)),
+		forward: forward,
+	}
+}
+
 func (p AllShortest) at(from, to int) (mid []int) {
 	return p.next[from+to*len(p.nodes)]
 }