Allow graph edges with custom keys

Author: Pierstoval

src/data-structures/graph/GraphEdge.js

@@ -2,22 +2,27 @@ export default class GraphEdge {
   /**
    * @param {GraphVertex} startVertex
    * @param {GraphVertex} endVertex
-   * @param {number} [weight=1]
+   * @param {number} [weight=0]
+   * @param key
    */
-  constructor(startVertex, endVertex, weight = 0) {
+  constructor(startVertex, endVertex, weight = 0, key = null) {
     this.startVertex = startVertex;
     this.endVertex = endVertex;
     this.weight = weight;
+    this.key = key;
   }
 
-  /**
-   * @return {string}
-   */
   getKey() {
+    if (this.key) {
+      return this.key;
+    }
+
     const startVertexKey = this.startVertex.getKey();
     const endVertexKey = this.endVertex.getKey();
 
-    return `${startVertexKey}_${endVertexKey}`;
+    this.key = `${startVertexKey}_${endVertexKey}`;
+
+    return this.key;
   }
 
   /**
@@ -35,6 +40,6 @@ export default class GraphEdge {
    * @return {string}
    */
   toString() {
-    return this.getKey();
+    return this.getKey().toString();
   }
 }

src/data-structures/graph/__test__/GraphEdge.test.js

@@ -7,8 +7,6 @@ describe('GraphEdge', () => {
     const endVertex = new GraphVertex('B');
     const edge = new GraphEdge(startVertex, endVertex);
 
-    expect(edge.getKey()).toBe('A_B');
-    expect(edge.toString()).toBe('A_B');
     expect(edge.startVertex).toEqual(startVertex);
     expect(edge.endVertex).toEqual(endVertex);
     expect(edge.weight).toEqual(0);
@@ -39,4 +37,18 @@ describe('GraphEdge', () => {
     expect(edge.endVertex).toEqual(vertexA);
     expect(edge.weight).toEqual(10);
   });
+
+  it('should return edges names as key', () => {
+    const edge = new GraphEdge(new GraphVertex('A'), new GraphVertex('B'), 0);
+
+    expect(edge.getKey()).toBe('A_B');
+    expect(edge.toString()).toBe('A_B');
+  });
+
+  it('should return custom key if defined', () => {
+    const edge = new GraphEdge(new GraphVertex('A'), new GraphVertex('B'), 0, 'custom_key');
+
+    expect(edge.getKey()).toEqual('custom_key');
+    expect(edge.toString()).toEqual('custom_key');
+  });
 });