Graph Traversal

graph.py

from typing import Iterable, List, Union, Set


class Graph:
    def __init__(self) -> None:

        # Adjacency set representation.
        # key: node name, value: adjacency node name set.
        self.edges = dict()

    def add_node(self, u: Union[int, str]) -> None:

        # Add a node.
        if u not in self.edges:
            self.edges[u] = set()

    def add_edge(self, u: Union[int, str], v: Union[int, str]) -> None:

        # Add an edge from u to v.
        self.add_node(u)
        self.add_node(v)
        self.edges[u].add(v)

    def get_neighbors(self, u: Union[int, str]) -> Set[Union[int, str]]:

        if u not in self.edges:
            raise RuntimeError(f"Node {u} does not exist in the graph.")

        return self.edges[u]

    def get_num_nodes(self) -> int:

        return len(self.edges)


class UndirectedGraph(Graph):
    def __init__(self) -> None:

        super(UndirectedGraph, self).__init__()

    def add_edge(self, u: Union[int, str], v: Union[int, str]) -> None:

        # Add an edge from u to v and from v to u.
        self.add_node(u)
        self.add_node(v)
        self.edges[u].add(v)
        self.edges[v].add(u)


class DirectedGraph(Graph):
    def __init__(self) -> None:

        super(DirectedGraph, self).__init__()

        # key: child name, value: parents set
        self.parents = dict()
        # key: parent name, value: children set
        self.children = self.edges

    def add_node(self, u: Union[int, str]) -> None:

        if u not in self.children:
            self.children[u] = set()
        if u not in self.parents:
            self.parents[u] = set()

    def add_edge(self, u: Union[int, str], v: Union[int, str]) -> None:

        self.add_node(u)
        self.add_node(v)
        self.children[u].add(v)
        self.parents[v].add(u)

    def get_parents(self, u: Union[int, str]) -> Set[Union[int, str]]:

        if u not in self.parents:
            raise RuntimeError(f"Node {u} does not exist in the graph.")

        return self.parents[u]

    def get_children(self, u: Union[int, str]) -> Set[Union[int, str]]:

        if u not in self.children:
            raise RuntimeError(f"Node {u} does not exist in the graph.")

        return self.children[u]


def create_dummy_directed_graph() -> DirectedGraph:

    graph = DirectedGraph()
    graph.add_edge(0, 1)
    graph.add_edge(0, 2)
    graph.add_edge(1, 2)
    graph.add_edge(2, 0)
    graph.add_edge(2, 3)
    graph.add_edge(3, 3)

    return graph


def create_dummy_undirected_graph() -> UndirectedGraph:

    graph = UndirectedGraph()
    graph.add_edge(1, 2)
    graph.add_edge(1, 3)
    graph.add_edge(2, 5)
    graph.add_edge(3, 5)
    graph.add_edge(2, 4)
    graph.add_edge(4, 5)
    graph.add_edge(4, 6)
    graph.add_edge(5, 6)

    return graph


if __name__ == "__main__":

    directed_graph = create_dummy_directed_graph()
    undirected_graph = create_dummy_undirected_graph()

search.py

# Graph search/traversal algorithms.
# These algorithms could be modified to achieve different goals for graphs.
from typing import Iterable, List, Union, Set
from queue import Queue
from graph import Graph, DirectedGraph, UndirectedGraph


def breadth_first_traversal(graph: Graph, start_nodes: List[int]) -> List[int]:

    # BFS supports multiple start nodes.
    # O(V + E) where V is the number of nodes and E is the number of edges.

    search_order = []

    num_nodes = graph.get_num_nodes()
    # First in first out (FIFO)
    q = Queue(maxsize=num_nodes)

    # Store the nodes that have been visited.
    visited = set()

    # Mark the source node as visited and enqueue it.
    for node in start_nodes:
        q.put(node)
        visited.add(node)

    while not q.empty():
        src_node = q.get()
        search_order.append(src_node)
        for tgt_node in graph.get_neighbors(src_node):
            if tgt_node not in visited:
                q.put(tgt_node)
                visited.add(tgt_node)

    return search_order


def recursive_depth_first_traversal(graph: Graph,
                                    start_node: int) -> List[int]:

    # Recursive algorithm is not favored for large graphs because of stack overflow.
    # O(V + E) where V is the number of nodes and E is the number of edges.

    def depth_first_recursion(graph: Graph, visited: Set, src_node: int,
                              search_order: List[int]):

        visited.add(src_node)
        search_order.append(src_node)
        for tgt_node in graph.get_neighbors(src_node):
            if tgt_node not in visited:
                depth_first_recursion(graph=graph,
                                      visited=visited,
                                      src_node=tgt_node,
                                      search_order=search_order)

    # Store the nodes that have been visited.
    visited = set()
    search_order = []

    if start_node not in visited:
        depth_first_recursion(graph=graph,
                              visited=visited,
                              src_node=start_node,
                              search_order=search_order)

    return search_order


def iterative_depth_first_traversal(graph: Graph,
                                    start_node: List[int]) -> List[int]:

    # O(V + E) where V is the number of nodes and E is the number of edges.

    # In Python, the list data structure can be used as stack.

    # Store the nodes that have been visited.
    visited = set()
    search_order = []
    # Last in first out (LIFO)
    stack = list()

    stack.append(start_node)
    while len(stack) != 0:
        src_node = stack.pop()
        if src_node not in visited:
            visited.add(src_node)
            search_order.append(src_node)
            for tgt_node in graph.get_neighbors(src_node):
                stack.append(tgt_node)

    return search_order


def create_dummy_directed_graph() -> DirectedGraph:

    graph = DirectedGraph()
    graph.add_edge(0, 1)
    graph.add_edge(0, 2)
    graph.add_edge(1, 2)
    graph.add_edge(2, 0)
    graph.add_edge(2, 3)
    graph.add_edge(3, 3)
    # Add an isolated node without connections.
    graph.add_node(4)

    return graph


def create_dummy_undirected_graph() -> UndirectedGraph:

    graph = UndirectedGraph()
    graph.add_edge(1, 2)
    graph.add_edge(1, 3)
    graph.add_edge(2, 5)
    graph.add_edge(3, 5)
    graph.add_edge(2, 4)
    graph.add_edge(4, 5)
    graph.add_edge(4, 6)
    graph.add_edge(5, 6)
    # Add an isolated node without connections.
    graph.add_node(0)

    return graph


if __name__ == "__main__":

    directed_graph = create_dummy_directed_graph()
    undirected_graph = create_dummy_undirected_graph()

    search_order = breadth_first_traversal(graph=directed_graph,
                                           start_nodes=[2, 0])
    print(search_order)
    search_order = breadth_first_traversal(graph=undirected_graph,
                                           start_nodes=[2, 0])
    print(search_order)

    search_order = recursive_depth_first_traversal(graph=directed_graph,
                                                   start_node=2)
    print(search_order)
    search_order = recursive_depth_first_traversal(graph=undirected_graph,
                                                   start_node=2)
    print(search_order)

    search_order = iterative_depth_first_traversal(graph=directed_graph,
                                                   start_node=2)
    print(search_order)
    search_order = iterative_depth_first_traversal(graph=undirected_graph,
                                                   start_node=2)
    print(search_order)