[Refactor] Graph now uses dict instead of defaultdict and supports all types as vertex names.

Author: exterkamp

structures/graph.py

@@ -1,4 +1,4 @@
-from collections import defaultdict
+from collections import defaultdict, deque
 
 class GraphNode():
 
@@ -18,7 +18,7 @@ class Graph():
     """
 
     def __init__(self, verticies):
-        self.graph = defaultdict(list)
+        self.graph = {}
         self.verticies = verticies
 
     def add_edge(self, source, destination):
@@ -33,7 +33,16 @@ def add_edge(self, source, destination):
             destination: the destination vertex
 
         """
-        self.graph[source].append(destination)
+        if len(self.graph) > self.verticies:
+            raise IndexError("Too many verticies in graph.")
+        
+        if source in self.graph:
+            self.graph[source].append(destination)
+        else:
+            self.graph[source] = [destination]
+        
+        if destination not in self.graph:
+            self.graph[destination] = []
 
     def has_cycle(self):
         """
@@ -96,21 +105,21 @@ def topological_sort(self):
 
         """
         visited = set()
-        stack = []
+        stack = deque()
 
         def dfs(vertex):
             visited.add(vertex)
-            for i in self.graph[vertex]:
-                if i not in visited:
-                    dfs(i)
+            for j in self.graph[vertex]:
+                if j not in visited:
+                    dfs(j)
 
-            stack.append(vertex)
-        
-        for i in range(self.verticies):
-            if i not in visited:
-                dfs(i)
+            stack.appendleft(vertex)
+
+        for key, _ in self.graph.items():
+            if key not in visited:
+                dfs(key)
 
-        return stack[::-1]
+        return list(stack)
 
 class WeightedGraphNode():
 

tests/test_graph_cycles.py

@@ -3,6 +3,17 @@
 
 class TestCycle(unittest.TestCase):
 
+    def test_non_int(self):
+
+        graph = Graph(4)
+
+        graph.add_edge('Tampa','Colorado')
+        graph.add_edge('Colorado','Las Vegas')
+        graph.add_edge('Las Vegas', 'Seattle')
+        graph.add_edge('Seattle','Anchorage')
+
+        self.assertEqual(['Tampa', 'Colorado', 'Las Vegas', 'Seattle', 'Anchorage'], graph.topological_sort())
+
     def test_cycles_true(self):
 
         graph = Graph(4)
@@ -59,4 +70,12 @@ def test_cycles_false(self):
         graph.add_edge(1,2)
         graph.add_edge(2,3)
 
-        self.assertEqual(False, graph.has_cycle())
+        self.assertEqual(False, graph.has_cycle())
+
+    def test_too_many_verticies_error(self):
+
+        graph = Graph(1)
+
+        graph.add_edge(0,1)
+        with self.assertRaises(IndexError):
+            graph.add_edge(1,2)

tests/test_topological_sort.py

@@ -15,7 +15,7 @@ def test_topological_sort(self):
 
         graph.add_edge(3,1)
 
-        self.assertEqual([5, 4, 2, 3, 1, 0], graph.topological_sort())
+        self.assertEqual([4, 5, 0, 2, 3, 1], graph.topological_sort())