Undirected graph class with adjacency matrix. DFS (recursive), DFS (iterative), BFS.

jwasham · jwasham · commit 4f106244e434 · 2016-08-02T20:19:04.000-07:00
diff --git a/graphs/undirected_graph_matrix.py b/graphs/undirected_graph_matrix.py
@@ -0,0 +1,147 @@
+import queue
+
+
+class UndirectedGraph(object):
+    """
+    Undirected Graph, with graph represented as an adjacency matrix
+    """
+
+    def __init__(self, vertices):
+        self.adjacency_matrix = []
+        for i in range(vertices):
+            self.adjacency_matrix.append([0] * vertices)
+
+    def add_edge(self, source, destination):
+        """
+        Adds an edge defined by vertices from source to destination
+        then destination to source in matrix
+        :param source:
+        :param destination:
+        :return:
+        """
+        self.adjacency_matrix[source][destination] = 1
+        self.adjacency_matrix[destination][source] = 1
+
+    def get_vertex(self):
+        """
+        Generator for returning the next vertex from the adjacency list
+        :return:
+        """
+        for v, __ in enumerate(self.adjacency_matrix):
+            yield v
+
+    def get_neighbor(self, vertex):
+        """
+        Generator for returning the next vertex adjacent to the given vertex
+        :param vertex:
+        :return:
+        """
+        for i, flag in enumerate(self.adjacency_matrix[vertex]):
+            if flag == 1:
+                yield i
+
+    def dfs_recursive(self):
+        """
+        Computes the parents for each vertex as determined through depth-first-search
+        (though not too meaningful in an undirected weightless graph)
+        :return: parents for each vertex
+        :rtype: dict
+        """
+        parents = {}
+
+        self.dfs_util(0, parents)
+
+        return parents
+
+    def dfs_util(self, vertex, parents):
+        for u in self.get_neighbor(vertex):
+            if u not in parents:
+                parents[u] = vertex
+                self.dfs_util(u, parents)
+
+    def dfs_iterative(self):
+        """
+        Computes the parents for each vertex as determined through depth-first-search
+        (though not too meaningful in an undirected weightless graph)
+        :return: parents for each vertex
+        :rtype: dict
+        """
+        parents = {}
+        to_visit = [0]
+
+        while to_visit:
+            v = to_visit.pop()
+
+            for neighbor in self.get_neighbor(v):
+                if neighbor not in parents:
+                    parents[neighbor] = v
+                    to_visit.append(neighbor)
+
+        return parents
+
+    def bfs(self):
+        """
+        Computes the the parents for each vertex as determined through breadth-first search
+        :return: parents for each vertex
+        :rtype: dict
+        """
+        parents = {}
+        to_visit = queue.Queue()
+        to_visit.put(0)
+
+        while not to_visit.empty():
+            v = to_visit.get()
+
+            for neighbor in self.get_neighbor(v):
+                if neighbor not in parents:
+                    parents[neighbor] = v
+                    to_visit.put(neighbor)
+
+        return parents
+
+
+def get_test_graph_1():
+    udg = UndirectedGraph(9)
+    udg.add_edge(0, 1)
+    udg.add_edge(1, 2)
+    udg.add_edge(2, 3)
+    udg.add_edge(1, 7)
+    udg.add_edge(3, 7)
+    udg.add_edge(7, 8)
+    udg.add_edge(3, 4)
+    udg.add_edge(3, 5)
+    udg.add_edge(4, 5)
+    udg.add_edge(5, 6)
+    udg.add_edge(6, 7)
+    udg.add_edge(6, 8)
+
+    return udg
+
+
+def test_dfs_recursive():
+    udg = get_test_graph_1()
+    p1 = udg.dfs_recursive()
+    assert(p1 == {0: 1, 1: 0, 2: 1, 3: 2, 4: 3, 5: 4, 6: 5, 7: 6, 8: 7})
+
+
+def test_dfs_iterative():
+    udg = get_test_graph_1()
+    p1 = udg.dfs_iterative()
+    assert(p1 == {0: 1, 1: 0, 2: 1, 3: 7, 4: 5, 5: 6, 6: 7, 7: 1, 8: 7})
+
+
+def test_bfs():
+    udg = get_test_graph_1()
+    p1 = udg.bfs()
+    assert(p1 == {0: 1, 1: 0, 2: 1, 3: 2, 4: 3, 5: 3, 6: 7, 7: 1, 8: 7})
+
+
+def main():
+    test_dfs_recursive()
+    test_dfs_iterative()
+    test_bfs()
+
+    print("Tests complete.")
+
+if __name__ == "__main__":
+    main()
