Graph implementation in Python

Intro

A graph is a data structure that consists of vertices that are connected via edges.

Using an adjacency matrix

# implement a graph using adjacency matrix
class Graph:
    def __init__(self, size):
        self.adj_matrix = []
        for i in range(size):
            self.adj_matrix.append([0 for i in range(size)])
        self.size = size

    def add_edge(self, v1, v2):
        if v1 != v2:
            self.adj_matrix[v1][v2] = 1
            self.adj_matrix[v2][v1] = 1

    def remove_edge(self, v1, v2):
        if self.adj_matrix[v1][v2] != 0:
            self.adj_matrix[v1][v2] = 0

    def print_matrix(self):
        for row in self.adj_matrix:
            for col in row:
                print(f"{col} ", end="")
            print()


g = Graph(5)
g.add_edge(0, 1)
g.add_edge(0, 2)
g.add_edge(1, 2)
g.add_edge(2, 0)
g.add_edge(2, 3)
g.print_matrix()

# Output:
# 0 1 1 0 0
# 1 0 1 0 0
# 1 1 0 1 0
# 0 0 1 0 0
# 0 0 0 0 0

Using adjacency matrix has a drawback to implement graph. Memory is allocated for all the edges whether it is present or not. This can be avoided by using the adjacency list.

Using an adjacency list

# Using linked list based deque to improve vertice removal efficiency

from collections import deque

class Graph:
    def __init__(self):
        self.adj_list = {}

    def add_edge(self, v1, v2):
        if v1 not in self.adj_list:
            self.adj_list[v1] = deque()

        if v2 not in self.adj_list[v1]:
            self.adj_list[v1].append(v2)

        if v2 not in self.adj_list:
            self.adj_list[v2] = deque()

        if v1 not in self.adj_list[v2]:
            self.adj_list[v2].append(v1)

    def remove_edge(self, v1, v2):
        if v1 in self.adj_list and v2 in self.adj_list:
            self.adj_list[v1].remove(v2)
            self.adj_list[v2].remove(v1)

    def print_graph(self):
        for v1 in self.adj_list:
            for v2 in self.adj_list[v1]:
                print(f"({v1},{v2}) ", end="")
            print()


g = Graph()
g.add_edge(0, 1)
g.add_edge(0, 2)
g.add_edge(1, 2)
g.add_edge(2, 0)
g.add_edge(2, 3)
g.print_graph()
print()
g.remove_edge(0, 1)
g.remove_edge(1, 2)
g.print_graph()

# Output:
# (0,1) (0,2)
# (1,0) (1,2)
# (2,0) (2,1) (2,3)
# (3,2)
#
# (0,2)
#
# (2,0) (2,3)
# (3,2)