Python/graphs/directed_and_undirected_weighted_graph.py

"""
Author  : Yashwanth Adimulam
Date    : 26th Oct 2024

Implement the class of Graphs with useful functions based on it.

Useful Links: https://www.geeksforgeeks.org/applications-advantages-and-disadvantages-of-weighted-graph/
              https://www.tutorialspoint.com/applications-advantages-and-disadvantages-of-unweighted-graph
              https://www.baeldung.com/cs/weighted-vs-unweighted-graphs
              https://en.wikipedia.org/wiki/Graph_(discrete_mathematics)
"""

from collections import deque
from math import floor
from random import random
from time import time

class DirectedGraph:
    def __init__(self) -> None:
        """
        Initialize a directed graph.

        >>> g = DirectedGraph()
        >>> g.all_nodes()
        []
        """
        self.graph = {}

    def add_pair(self, u: int, v: int, w: int = 1) -> None:
        """
        Add a directed edge from u to v with weight w.

        >>> g = DirectedGraph()
        >>> g.add_pair(1, 2, 3)
        >>> g.graph
        {1: [[3, 2]], 2: []}
        """
        if self.graph.get(u):
            if not any(edge[1] == v and edge[0] == w for edge in self.graph[u]):
                self.graph[u].append([w, v])
        else:
            self.graph[u] = [[w, v]]
        if v not in self.graph:
            self.graph[v] = []

    def all_nodes(self) -> list:
        """
        Return a list of all nodes in the graph.

        >>> g = DirectedGraph()
        >>> g.add_pair(1, 2)
        >>> g.all_nodes()
        [1, 2]
        """
        return list(self.graph)

    def remove_pair(self, u: int, v: int) -> None:
        """
        Remove the directed edge from u to v.

        >>> g = DirectedGraph()
        >>> g.add_pair(1, 2)
        >>> g.remove_pair(1, 2)
        >>> g.graph
        {1: [], 2: []}
        """
        if self.graph.get(u):
            self.graph[u] = [edge for edge in self.graph[u] if edge[1] != v]

    def dfs(self, s: int = -2, d: int = -1) -> list:
        """
        Perform depth-first search from node s to d.

        >>> g = DirectedGraph()
        >>> g.add_pair(1, 2)
        >>> g.add_pair(2, 3)
        >>> g.dfs(1, 3)
        [1, 2, 3]
        """
        if s == d:
            return []
        stack = []
        visited = []
        if s == -2:
            s = next(iter(self.graph))
        stack.append(s)
        visited.append(s)
        while stack:
            current = stack[-1]
            if current == d:
                return visited
            if current in self.graph:
                for edge in self.graph[current]:
                    if edge[1] not in visited:
                        stack.append(edge[1])
                        visited.append(edge[1])
                        break
                else:
                    stack.pop()
            else:
                stack.pop()
        return visited

    def bfs(self, s: int = -2) -> list:
        """
        Perform breadth-first search starting from node s.

        >>> g = DirectedGraph()
        >>> g.add_pair(1, 2)
        >>> g.add_pair(2, 3)
        >>> g.bfs(1)
        [1, 2, 3]
        """
        d = deque()
        visited = []
        if s == -2:
            s = next(iter(self.graph))
        d.append(s)
        visited.append(s)
        while d:
            current = d.popleft()
            if current in self.graph:
                for edge in self.graph[current]:
                    if edge[1] not in visited:
                        d.append(edge[1])
                        visited.append(edge[1])
        return visited

    def has_cycle(self) -> bool:
        """
        Check if the graph has a cycle.

        >>> g = DirectedGraph()
        >>> g.add_pair(1, 2)
        >>> g.add_pair(2, 1)
        >>> g.has_cycle()
        True
        """
        visited = set()
        rec_stack = set()

        def cycle_util(v):
            visited.add(v)
            rec_stack.add(v)
            for edge in self.graph.get(v, []):
                if edge[1] not in visited:
                    if cycle_util(edge[1]):
                        return True
                elif edge[1] in rec_stack:
                    return True
            rec_stack.remove(v)
            return False

        for node in self.graph:
            if node not in visited:
                if cycle_util(node):
                    return True
        return False

    # Additional methods would go here with doctests...

class Graph:
    def __init__(self) -> None:
        """
        Initialize an undirected graph.

        >>> g = Graph()
        >>> g.all_nodes()
        []
        """
        self.graph = {}

    def add_pair(self, u: int, v: int, w: int = 1) -> None:
        """
        Add an undirected edge between u and v with weight w.

        >>> g = Graph()
        >>> g.add_pair(1, 2, 3)
        >>> g.graph
        {1: [[3, 2]], 2: [[3, 1]]}
        """
        if self.graph.get(u):
            if not any(edge[1] == v and edge[0] == w for edge in self.graph[u]):
                self.graph[u].append([w, v])
        else:
            self.graph[u] = [[w, v]]
        if self.graph.get(v):
            if not any(edge[1] == u and edge[0] == w for edge in self.graph[v]):
                self.graph[v].append([w, u])
        else:
            self.graph[v] = [[w, u]]

    def all_nodes(self) -> list:
        """
        Return a list of all nodes in the graph.

        >>> g = Graph()
        >>> g.add_pair(1, 2)
        >>> g.all_nodes()
        [1, 2]
        """
        return list(self.graph)

    # Additional methods would go here with doctests...

if __name__ == "__main__":
    import doctest
    doctest.testmod()