Merge pull request #15 from datavis-tech/weighted-edges

adding edge weights

Author: curran

.gitignore

@@ -1 +1,3 @@
 node_modules
+package-lock.json
+*.swp

README.md

@@ -116,14 +116,26 @@ Removes the specified node. Returns *graph* to support method chaining. The argu
 
 ### Adding and Removing Edges
 
-<a name="add-edge" href="#add-edge">#</a> <i>graph</i>.<b>addEdge</b>(<i>u</i>, <i>v</i>)
+<a name="add-edge" href="#add-edge">#</a> <i>graph</i>.<b>addEdge</b>(<i>u</i>, <i>v</i>[,<i>weight</i>])
 
 Adds an edge from node *u* to node *v*. Returns *graph* to support method chaining. The arguments *u* and *v* are string identifiers for nodes. This function also adds *u* and *v* as nodes if they were not already added.
 
+The last argument *weight* (optional) specifies the weight of this edge.
+
 <a name="remove-edge" href="#remove-edge">#</a> <i>graph</i>.<b>removeEdge</b>(<i>u</i>, <i>v</i>)
 
 Removes the edge from node *u* to node *v*. Returns *graph* to support method chaining. The arguments *u* and *v* are string identifiers for nodes. This function does not remove the nodes *u* and *v*. Does nothing if the edge does not exist.
 
+### Working with Edge Weights
+
+<a name="set-edge-weight" href="#set-edge-weight">#</a> <i>graph</i>.<b>setEdgeWeight</b>(<i>u</i>, <i>v</i>, <i>weight</i>)
+
+Sets the *weight* (a number) of the edge from node *u* to node *v*.
+
+<a name="get-edge-weight" href="#get-edge-weight">#</a> <i>graph</i>.<b>getEdgeWeight</b>(<i>u</i>, <i>v</i>, <i>weight</i>)
+
+Gets the *weight* of the edge from node *u* to node *v*. If no weight was previously set on this edge, then the value 1 is returned.
+
 ### Querying the Graph
 
 <a name="nodes" href="#nodes">#</a> <i>graph</i>.<b>nodes</b>()

index.js

@@ -9,6 +9,8 @@ module.exports = function Graph(serialized){
     adjacent: adjacent,
     addEdge: addEdge,
     removeEdge: removeEdge,
+    setEdgeWeight: setEdgeWeight,
+    getEdgeWeight: getEdgeWeight,
     indegree: indegree,
     outdegree: outdegree,
     depthFirstSearch: depthFirstSearch,
@@ -22,6 +24,11 @@ module.exports = function Graph(serialized){
   // Values are adjacent node id arrays.
   var edges = {};
 
+  // The weights of edges.
+  // Keys are string encodings of edges.
+  // Values are weights (numbers).
+  var edgeWeights = {};
+
   // If a serialized graph was passed into the constructor, deserialize it.
   if(serialized){
     deserialize(serialized);
@@ -72,12 +79,36 @@ module.exports = function Graph(serialized){
     return edges[node] || [];
   }
 
+  // Computes a string encoding of an edge,
+  // for use as a key in an object.
+  function encodeEdge(u, v){
+    return u + "|" + v;
+  }
+
+  // Sets the weight of the given edge.
+  function setEdgeWeight(u, v, weight){
+    edgeWeights[encodeEdge(u, v)] = weight;
+    return graph;
+  }
+
+  // Gets the weight of the given edge.
+  // Returns 1 if no weight was previously set.
+  function getEdgeWeight(u, v){
+    var weight = edgeWeights[encodeEdge(u, v)];
+    return weight === undefined ? 1 : weight;
+  }
+
   // Adds an edge from node u to node v.
   // Implicitly adds the nodes if they were not already added.
-  function addEdge(u, v){
+  function addEdge(u, v, weight){
     addNode(u);
     addNode(v);
     adjacent(u).push(v);
+
+    if (weight !== undefined) {
+      setEdgeWeight(u, v, weight);
+    }
+
     return graph;
   }
 

test.js

@@ -311,6 +311,27 @@ describe("Graph", function() {
       checkSerialized(graph.serialize());
     });
   });
+
+  describe("Edge Weights", function() {
+
+    it("Should set and get an edge weight.", function (){
+      var graph = Graph().addEdge("a", "b", 5);
+      assert.equal(graph.getEdgeWeight("a", "b"), 5);
+    });
+
+    it("Should set edge weight via setEdgeWeight.", function (){
+      var graph = Graph()
+        .addEdge("a", "b")
+        .setEdgeWeight("a", "b", 5);
+      assert.equal(graph.getEdgeWeight("a", "b"), 5);
+    });
+
+    it("Should return weight of 1 if no weight set.", function (){
+      var graph = Graph().addEdge("a", "b");
+      assert.equal(graph.getEdgeWeight("a", "b"), 1);
+    });
+
+  });
 });
 
 function contains(arr, item){