Implement undirected adjacency list with unit tests

skirkpatrick · skirkpatrick · commit 109a592559ae · 2013-04-01T14:06:18.000-07:00
diff --git a/algorithms/data_structures/__init__.py b/algorithms/data_structures/__init__.py
diff --git a/algorithms/data_structures/undirected_adjacency_matrix.py b/algorithms/data_structures/undirected_adjacency_matrix.py
@@ -0,0 +1,58 @@
+"""
+    undirected_adjacency_matrix.py
+
+    Implementation of a graph as an undirected adjacency matrix.
+
+    Undirected Adjacency Matrix Overview:
+    --------------------------------------
+    Edges are stored as a two-dimensional matrix of weights.
+    Vertices are represented as the rows and columns of the matrix.
+    Since the graph is undirected, for any two vertices x and y:
+        edge(x,y) == edge(y,x)
+
+
+    Time complexity:
+    ----------------
+    V = Number of vertices
+
+    Storage:       O(V**2)
+    Add vertex:    O(V**2)
+    Add edge:      O(1)
+    Remove vertex: O(v**2)
+    Remove edge:   O(1)
+    Query:         O(1)
+
+    Graph comparison: http://en.wikipedia.org/wiki/Graph_(abstract_data_type)
+    Wiki: http://en.wikipedia.org/wiki/Adjacency_matrix
+
+"""
+
+
+class UndirectedAdjacencyMatrix:
+    def __init__(self):
+        self.edge = {}
+        self.numVertices = 0
+
+    def addVertex(self, key):
+        if key in self.edge:
+            return
+        self.edge[key] = {v: float('inf') for v in self.edge}
+        for i in self.edge:
+            self.edge[i][key] = float('inf')
+        self.numVertices += 1
+
+    def removeVertex(self, key):
+        if key not in self.edge:
+            return
+        del self.edge[key]
+        for i in self.edge:
+            del self.edge[i][key]
+        self.numVertices -= 1
+
+    def addEdge(self, x, y, weight = 1):
+        if all (v in self.edge for v in (x, y)):
+            self.edge[x][y] = self.edge[y][x] = weight
+
+    def removeEdge(self, x, y):
+        if all (v in self.edge for v in (x, y)):
+            self.edge[x][y] = self.edge[y][x] = float('inf')
diff --git a/algorithms/tests/test_data_structures.py b/algorithms/tests/test_data_structures.py
@@ -0,0 +1,75 @@
+""" Unit Tests for data structures """
+import unittest
+from ..data_structures import undirected_adjacency_matrix
+
+
+class TestUndirectedAdjacencyMatrix(unittest.TestCase):
+    """
+    Tests edge cases for Undirected Adjacency Matrix
+    """
+
+    def setUp(self):
+        self.graph = undirected_adjacency_matrix.UndirectedAdjacencyMatrix()
+        self.expected = []
+
+        #Test empty graph
+        self.expected.append({})
+
+        #Test adding vertices
+        self.expected.append({v: {i: float('inf') for i in range(20)}
+                                 for v in range(20)})
+
+        #Test removing vertices
+        self.expected.append({v: {i: float('inf') for i in range(0, 20, 2)}
+                                 for v in range(0, 20, 2)})
+
+        #Test adding edges
+        self.expected.append({v: {i: 1 for i in range(0, 20, 2)}
+                                 for v in range(0, 20, 2)})
+
+        #Test removing edges
+        self.expected.append({v: {i: 1 if (i < 10 and v < 10) or
+                                     (i >= 10 and v >= 10) else float('inf')
+                                     for i in range(0, 20, 2)}
+                                 for v in range(0, 20, 2)})
+
+        #Test empty graph after removing all vertices
+        self.expected.append({})
+
+    def test_undirected_adjacency_matrix(self):
+        #Test empty graph
+        self.assertEqual(self.expected[0], self.graph.edge)
+
+        #Test adding vertices
+        #Should result in 20 vertices with no edges
+        for i in range(20):
+            self.graph.addVertex(i)
+        self.assertEqual(self.expected[1], self.graph.edge)
+
+        #Test removing vertices
+        #Should result in 10 even-numbered vertices with no edges
+        for i in range(1, 20, 2):
+            self.graph.removeVertex(i)
+        self.assertEqual(self.expected[2], self.graph.edge)
+
+        #Test adding edges
+        #Should result in a complete graph
+        for i in range(0, 20, 2):
+            for j in range(0, 20, 2):
+                self.graph.addEdge(i, j)
+        self.assertEqual(self.expected[3], self.graph.edge)
+
+        #Test removing edges
+        #Should result in two complete components: vertices 0-8, 10-18
+        for i in range(0, 10, 2):
+            for j in range(10, 20, 2):
+                self.graph.removeEdge(i ,j)
+        for i in range(10, 20, 2):
+            for j in range(0, 10, 2):
+                self.graph.removeEdge(i, j)
+        self.assertEqual(self.expected[4], self.graph.edge)
+
+        #Test empty graph after removing all vertices
+        for i in range(0, 20, 2):
+            self.graph.removeVertex(i)
+        self.assertEqual(self.expected[5], self.graph.edge)
