TheAlgorithms · raklaptudirm · Oct 29, 2021 · Oct 25, 2021 · Oct 28, 2021 · Oct 29, 2021
@@ -1,37 +1,54 @@
-// This file contains the simple structural implementation of undirected
-// graph. Used in coloring algorithms [here](./coloring)
-// Author(s): [Shivam](https://github.com/Shivam010)
+// This file contains the simple structural implementation of
+// directed & undirected graphs used within the graph package
+// Author(s): [Shivam](https://github.com/Shivam010), [Tahmeed](https://github.com/Tahmeed156)
 
 package graph
 
-// UndirectedGraph provides a structure to store an undirected graph.
+// Graph provides a structure to store the graph.
 // It is safe to use its empty object.
-type UndirectedGraph struct {
+type Graph struct {
 	vertices int
-	edges    map[int]map[int]struct{}
+	edges    map[int]map[int]int // Stores weight of an edge
+	Directed bool                // Differentiate directed/undirected graphs
 }
 
-// AddVertex will add a new vertex in the graph, if the vertex already
-// exist it will do nothing
-func (g *UndirectedGraph) AddVertex(v int) {
+// Constructor functions for graphs (undirected by default)
+func New(v int) *Graph {
+	return &Graph{
+		vertices: v,
+	}
+}
+
+// AddVertex will add a new vertex in the graph.
+// If the vertex already exists it will do nothing.
+func (g *Graph) AddVertex(v int) {
 	if g.edges == nil {
-		g.edges = make(map[int]map[int]struct{})
+		g.edges = make(map[int]map[int]int)
 	}
 
 	// Check if vertex is present or not
 	if _, ok := g.edges[v]; !ok {
-		g.vertices++
-		g.edges[v] = make(map[int]struct{})
+		g.edges[v] = make(map[int]int)
 	}
 }
 
 // AddEdge will add a new edge between the provided vertices in the graph
-func (g *UndirectedGraph) AddEdge(one, two int) {
+func (g *Graph) AddEdge(one, two int) {
+	// Add an edge with 0 weight
+	g.AddWeightedEdge(one, two, 0)
+}
+
+// AddWeightedEdge will add a new weighted edge between the provided vertices in the graph
+func (g *Graph) AddWeightedEdge(one, two, weight int) {
 	// Add vertices: one and two to the graph if they are not present
 	g.AddVertex(one)
 	g.AddVertex(two)
 
-	// and finally add the edges: one->two and two->one for undirected graph
-	g.edges[one][two] = struct{}{}
-	g.edges[two][one] = struct{}{}
+	// And finally add the edges
+	// one->two and two->one for undirected graph
+	// one->two for directed graphs
+	g.edges[one][two] = weight
+	if !g.Directed {
+		g.edges[two][one] = weight
+	}
 }
@@ -1,3 +1,131 @@
-// Empty test file to keep track of all the tests for the algorithms.
+// Tests for directed and undirected graphs
 
 package graph
+
+import (
+	"fmt"
+	"testing"
+)
+
+var graphTestCases = []struct {
+	name     string
+	edges    [][]int
+	vertices int
+}{
+	{
+		"single edge",
+		[][]int{
+			{0, 1, 1},
+		},
+		2,
+	},
+	{
+		"many edges",
+		[][]int{
+			{0, 1, 1},
+			{0, 2, 2},
+			{1, 3, 4},
+			{3, 4, 3},
+			{4, 8, 3},
+			{4, 9, 1},
+			{7, 8, 2},
+			{8, 9, 2},
+		},
+		10,
+	},
+	{
+		"cycles",
+		[][]int{
+			{0, 1, 1},
+			{0, 2, 2},
+			{1, 3, 4},
+			{3, 4, 3},
+			{4, 2, 1},
+		},
+		5,
+	},
+	{
+		"disconnected graphs",
+		[][]int{
+			{0, 1, 5},
+			{2, 4, 5},
+			{3, 8, 5},
+		},
+		2,
+	},
+}
+
+func TestDirectedGraph(t *testing.T) {
+
+	// Testing self-loops separately only for directed graphs.
+	// For undirected graphs each edge already creates a self-loop.
+	directedGraphTestCases := append(graphTestCases, struct {
+		name     string
+		edges    [][]int
+		vertices int
+	}{
+		"self-loops",
+		[][]int{
+			{0, 1, 1},
+			{1, 2, 2},
+			{2, 1, 3},
+		},
+		3,
+	})
+
+	for _, test := range directedGraphTestCases {
+		t.Run(fmt.Sprint(test.name), func(t *testing.T) {
+			// Initializing graph, adding edges
+			graph := New(test.vertices)
+			graph.Directed = true
+			for _, edge := range test.edges {
+				graph.AddWeightedEdge(edge[0], edge[1], edge[2])
+			}
+
+			if graph.vertices != test.vertices {
+				t.Errorf("Number of vertices, Expected: %d, Computed: %d", test.vertices, graph.vertices)
+			}
+			edgeCount := 0
+			for _, e := range graph.edges {
+				edgeCount += len(e)
+			}
+			if edgeCount != len(test.edges) {
+				t.Errorf("Number of edges, Expected: %d, Computed: %d", len(test.edges), edgeCount)
+			}
+			for _, edge := range test.edges {
+				if val, found := graph.edges[edge[0]][edge[1]]; !found || val != edge[2] {
+					t.Errorf("Edge {%d->%d (%d)} not found", edge[0], edge[1], edge[2])
+				}
+			}
+		})
+	}
+}
+
+func TestUndirectedGraph(t *testing.T) {
+
+	for _, test := range graphTestCases {
+		t.Run(fmt.Sprint(test.name), func(t *testing.T) {
+			// Initializing graph, adding edges
+			graph := New(test.vertices)
+			for _, edge := range test.edges {
+				graph.AddWeightedEdge(edge[0], edge[1], edge[2])
+			}
+
+			if graph.vertices != test.vertices {
+				t.Errorf("Number of vertices, Expected: %d, Computed: %d", test.vertices, graph.vertices)
+			}
+			edgeCount := 0
+			for _, e := range graph.edges {
+				edgeCount += len(e)
+			}
+			if edgeCount != len(test.edges)*2 {
+				t.Errorf("Number of edges, Expected: %d, Computed: %d", len(test.edges)*2, edgeCount)
+			}
+			for _, edge := range test.edges {
+				if val, found := graph.edges[edge[0]][edge[1]]; !found || val != edge[2] {
+					t.Errorf("Edge {%d->%d (%d)} not found", edge[0], edge[1], edge[2])
+				}
+			}
+		})
+	}
+}