Matrix implementation of graph using list and dict.

matrix_graph.py

import unittest


class ListGraph(object):
    def __init__(self, number_of_vertices):
        self.number_of_vertices = number_of_vertices
        self.matrix = [[0] * number_of_vertices for _ in range(number_of_vertices)]

    def add_edge(self, v1, v2):
        self.matrix[v1][v2] = 1
        self.matrix[v2][v1] = 1

    def add_vertex(self):
        self.number_of_vertices += 1
        for vertices in self.matrix:
            vertices.append(0)
        self.matrix.append([0] * self.number_of_vertices)

    def remove_edge(self, v1, v2):
        self.matrix[v1][v2] = 0
        self.matrix[v2][v1] = 0

    def remove_vertex(self, v):
        for vertices in self.matrix:
            vertices.remove(vertices[v])
        self.matrix.remove(self.matrix[v])
        self.number_of_vertices -= 1



class DictGraph(dict):
    def __init__(self, vertices):
        self.vertices = vertices

    def add_edge(self, v1, v2):
        if v1 in self.vertices and v2 in self.vertices:
            self[(v1, v2)] = 1
            self[(v2, v1)] = 1

    def add_vertex(self, v):
        self.vertices.append(v)

    def remove_edge(self, v1, v2):
        del self[(v1, v2)]
        del self[(v2, v1)]

    def remove_vertex(self, v):
        keys = list(self.keys())
        for v1, v2 in keys:
            if v1 == v or v2 == v and self.get((v1, v2)):
                del self[(v1, v2)]
        self.vertices.remove(v)


class ListGraphTest(unittest.TestCase):

    def test_add_edge_for_existing_vertices(self):
        number_of_vertices = 3
        graph = ListGraph(number_of_vertices)
        graph.add_edge(2, 1)
        self.assertEquals(1, graph.matrix[2][1])
        self.assertEquals(1, graph.matrix[1][2])

    def test_add_vertex(self):
        number_of_vertices = 3
        graph = ListGraph(number_of_vertices)
        expected_number_of_vertices = number_of_vertices + 1
        graph.add_vertex()

        self.assertEquals(expected_number_of_vertices, graph.number_of_vertices)
        self.assertEquals(expected_number_of_vertices, len(graph.matrix))
        for vertices in graph.matrix:
            self.assertEquals(expected_number_of_vertices, len(vertices))

    def test_remove_existing_edge(self):
        graph = ListGraph(3)
        graph.add_edge(0, 1)
        graph.add_edge(2, 1)

        graph.remove_edge(2, 1)
        self.assertEquals(1, graph.matrix[0][1])
        self.assertEquals(0, graph.matrix[2][1])
        self.assertEquals(0, graph.matrix[1][2])

    def test_remove_vertex(self):
        number_of_vertices = 3
        graph = ListGraph(3)
        expected_number_of_vertices = number_of_vertices - 1
        graph.add_edge(0, 1)
        graph.add_edge(2, 1)

        graph.remove_vertex(2)
        self.assertEquals(expected_number_of_vertices, graph.number_of_vertices)
        self.assertEquals(expected_number_of_vertices, len(graph.matrix))
        for vertices in graph.matrix:
            self.assertEquals(expected_number_of_vertices, len(vertices))
        self.assertEquals(1, graph.matrix[0][1])


class DictGraphTest(unittest.TestCase):

    def test_add_edge_for_existing_vertices(self):
        graph = DictGraph([1, 2, 3])
        graph.add_edge(1, 2)

        self.assertIsNotNone(graph[(1, 2)])
        self.assertIsNotNone(graph[(2, 1)])

    def test_add_edge_for_not_existing_vertices(self):
        graph = DictGraph([1, 2, 3])
        graph.add_edge(4, 2)

        self.assertIsNone(graph.get((4, 2)))
        self.assertIsNone(graph.get((2, 4)))

    def test_remove_edge(self):
        graph = DictGraph([1, 2, 3])
        graph.add_edge(1, 2)
        graph.add_edge(3, 2)
        graph.remove_edge(3, 2)

        self.assertIsNotNone(graph.get((1, 2)))
        self.assertIsNone(graph.get((3, 2)))
        self.assertIsNone(graph.get((2, 3)))

    def test_remove_vertex(self):
        graph = DictGraph([1, 2, 3])
        graph.add_edge(1, 2)
        graph.add_edge(1, 3)

        graph.remove_vertex(3)
        self.assertNotIn(3, graph.vertices)
        self.assertNotIn((1, 3), graph.keys())
        self.assertIn((1, 2), graph.keys())