graphmatrix.py

'''
  Graph represented as a matrix
'''
class GraphMatrix:

    '''
        Constructor, initializes the graph with a list of vertices
        and the type of graph (directed or undirected, weighted or unweighted)
    '''
    def __init__(self, vertexList = [], directed = False, weighted = False):
        self.vertexList = vertexList
        self.directed = directed
        self.weighted = weighted

        length = len(vertexList)
        self.matrix =  [[0] * length for _ in range(length)]

    '''
        Add an edge to the graph
    '''
    def addEdge(self, start, end, weight = 1):
        if self.matrix[start][end] != 0:
            return

        edgeWeight = weight if self.weighted else 1
        self.matrix[start][end] = edgeWeight
        if not self.directed:
            self.matrix[end][start] = edgeWeight

    '''
        Remove an edge from the graph
    '''
    def removeEdge(self, start,end):
        if self.matrix[start][end] == 0:
            return

        self.matrix[start][end] = 0
        if not self.directed:
            self.matrix[end][start] = 0

    '''
        Add a vertex to the graph
    '''
    def addVertex(self, vertex):
        length = len(self.vertexList)
        for row in self.matrix:
            row.append(0)
        self.matrix.append([0] * (length + 1))
        self.vertexList.append(vertex)
    
    '''
        Remove a vertex from the graph
    '''
    def removeVertex(self, vertex):
        index = self.vertexList.index(vertex)
        if(index == -1):
            return
        self.vertexList.remove(vertex)
        self.matrix.pop(index)
        for row in self.matrix:
            row.pop(index)