Graph Algorithms

def _build_graph(edges: list[list[int]], num_nodes: int) -> list[set]:
    """Builds a adjacency-list type graph from a list of edges.
    :param edges: Each item is an edge [i, j] such that we have an edge i -> j.

    Sample input::
        edges=[[1,0], [2,0], [3,2], [1,3]], num_nodes=5
        [1] -> [0]   [4]
         |      ^
         v      |
        [3] -> [2]

    Sample output::
        graph=[set(), {0,3}, {0}, {2}, set()]
    """
    adj_set = [set() for _ in range(num_nodes)]
    for tail, head in edges:
        adj_set[tail].add(head)
    print(f"{adj_set=}")
    return adj_set

@dataclass
class Edge:
    nodes: tuple['Node', 'Node']

@dataclass
class Node:
    val: int = 0
    edges: Optional[set['Edge']] = None

"""
Graph
        (0) -- (1)
         |
        (2)
"""
nodes = [Node(val=0), Node(val=1), Node(val=2)]
edges = [Edge(nodes=(nodes[0], nodes[1])), Edge(nodes=(node[0], node[2]))]
nodes[0].edges = [edge[0], edge[1]]
nodes[1].edges = [edge[0]]
nodes[2].edges = [edge[1]]

@dataclass
class Edge:
    tail: 'Node'
    head: 'Node'

@dataclass
class Node:
    val: int = 0
    outward_edges: Optional[set['Edge']] = None
    # This variable is optional and adds to storage. We
    # can get away with just the outward edges.
    inward_edges: Optional[set['Edge']] = None

"""
Graph
        (0) --> (1)
         Ʌ
         |
        (2)
"""
nodes = [Node(val=0), Node(val=1), Node[val=2]]
edges = [
    Edge(tail=nodes[0], head=nodes[1]),
    Edge(tail=nodes[2], head=nodes[0]),
]
nodes[0].outward_edges = {edges[0]}
nodes[0].inward_edges = {edges[1]}
nodes[1].outward_edges = {}
nodes[1].inward_edges = {edges[0]}
nodes[2].outward_edges = {edges[1]}
nodes[2].inward_edges = {edges[0]}

import typing
from dataclasses import dataclass
# Graph represented as an adjacency list, `len(graph)` is num nodes and
# `len(graph[i])` is number of adjacent numbers from node `i`.
Graph = typing.NewType('Graph', list[list[int]])

@dataclass
class Index:
    ix: int

def topsort(graph: Graph) -> list[int]:
    if not graph:
        return []
    n = len(graph)
    output, output_ix, seen = [None]*n,  Index(ix=n-1), set()
    for i in range(n):
        if i in seen:
            continue
        _dfs(graph=graph, node=i, seen=seen, output=output, output_ix=output_ix)
    return output

def _dfs(
    graph: Graph, node: int, seen: set[int], output: list[int], output_ix: Index,
) -> None:
    seen.add(node)
    # This will be only be empty for the sink node; equally this loop will be a no-op
    # if all the descendents have already been visitied
    for adj in graph[node]:
        if adj in seen:
            continue
        _dfs(graph=graph, node=adj, seen=seen, output=output, output_ix=output_ix)
    output[output_ix.ix] = node
    output_ix.ix -= 1

#          ┌────┐
#   ┌──────► 1  ├───────┐
#   │      └────┘       │
# ┌─┴──┐             ┌──▼─┐
# │ 0  │             │ 3  │
# └─┬──┘             └──▲─┘
#   │      ┌────┐       │
#   └──────► 2  ├───────┘
#          └────┘
print(topsort(graph=Graph([[1,2],[3], [3], []])))

def has_dag_cycle(graph: list[set]) -> bool:
    """Checks if a DAG has a cycle.

    :param graph: Each item in a list is a node and each element in the set are the
        nodes at the head of the outward arc from this node.

    Sample cyclic input::
        [{2}, {0}, {3}, {1}]

        [1] ->  [0]
         ^       |
         |       v
        [3] <-  [2]

    Sample non-cyclic input::
        [{}, {0,3}, {0}, {2}, {}]
          0    1     2    3    4

        [1] -> [0]     [4]
        |       ^
        v       |
        [3] -> [2]
    """
    explore = set(range(len(graph)))
    visiting, done = set(), set()
    while explore:
        if _cycle_dfs(
            node=next(iter(explore)),
            graph=graph,
            explore=explore,
            visiting=visiting,
            done=done,
        ):
            return True
    else:
        return False

def _cycle_dfs(node, graph, explore, visiting, done):
    _move_node(node, explore, visiting)
    for adj in graph[node]:
        if adj in done:
            continue
        if adj in visiting:
            # Found a cycle. This node was visited earlier while recursing to this node.
            # `adj` is both an ancestor and and a descendent of the current node.
            return True
        if _cycle_dfs(
            node=adj, graph=graph, explore=explore, visiting=visiting, done=done,
        ):
            return True
    else:
        _move_node(node, visiting, done)
        return False

def _move_node(node, from_set, to_set) -> None:
    """Side effect, mutates the input sets"""
    from_set.discard(node)
    to_set.add(node)

from collections.abc import Iterable
from dataclasses import dataclass
from dataclasses import field
from typing import Generator
from typing import Hashable

@dataclass
class DSU:
    """Disjoint Set Union implementation."""

    parent: dict[Hashable, Hashable] = field(default_factory=dict)
    rank: dict[Hashable, int] = field(default_factory=dict)

    def add_nodes(self, ns: Iterable[Hashable]) -> 'DSU':
        for n in ns:
            self.add_node(n=n)
        return self

    def add_node(self, n: Hashable) -> 'DSU':
        self.parent[n] = n
        self.rank[n] = 1
        return self

    def find_leader(self, n: Hashable) -> Hashable:
        if self.parent[n] == n:
            return n
        self.parent[n] = self.find_leader(n=self.parent[n])
        return self.parent[n]

    def union(self, n1: Hashable, n2: Hashable) -> 'DSU':
        n1_leader = self.find_leader(n=n1)
        n2_leader = self.find_leader(n=n2)
        if n1_leader == n2_leader:
            return
        if self.rank[n1_leader] > self.rank[n2_leader]:
            self.parent[n2_leader] = n1_leader
        elif self.rank[n1_leader] < self.rank[n2_leader]:
            self.parent[n1_leader] = n2_leader
        else:
            # equal ht trees
            self.parent[n2_leader] = n1_leader
            self.rank[n1_leader] += 1
        return self

    def is_connected(self, n1: Hashable, n2: Hashable) -> bool:
        return self.find_leader(n=n1) == self.find_leader(n=n2)

def has_undirected_cycle(graph: dict[str, set[str]]) -> bool:
    """Finds cycles in an undirected graph.

    :param graph: The keys of the dict are the graph nodes which are strings. The set
        of strings they hold are the set of nodes connected to them.
    """
    dsu = DSU()
    dsu.add_nodes(ns=graph.keys())
    for n1, n2 in _iter_edges(graph):
        if dsu.is_connected(n1, n2):
            return True
        dsu.union(n1, n2)
    else:
        return False

def _iter_edges(graph: dict[str, set[str]]) -> Generator[tuple, None, None]:
    """Iterate through edges in a undirected graph. In an undirected graph, two
    connected nodes A and B will have each other in their edge set and hence the edge
    appears twice, we have to ensure we only yield an edge once.
    """
    seen = set()
    for n, adjs in graph.items():
        for adj in adjs:
            edge = tuple(sorted([n, adj]))
            if edge in seen:
                continue
            seen.add(edge)
            yield edge

# [d-a-c] [f] [b-e]
graph = {'a': {'c', 'd'}, 'b': {'e'}, 'c': {'a'}, 'd': {'a'}, 'e': {'b'}, 'f': {}}
has_undirected_cycle(graph)
# False

# [f]
#
#  b - a - c
#  | /  \
#  e     d
graph = {
    'a': {'c', 'd', 'e', 'b'}, 'b': {'e', 'a'}, 'c': {'a'},
    'd': {'a'}, 'e': {'b', 'a'}, 'f': {}
}
has_undirected_cycle(graph)
# True

def has_undirected_cycle_dfs(graph: dict[str, set[str]]) -> bool:
    seen = set()
    # The graph may not be fully connected, we need to DFS through each connected set.
    for n in graph:
        if n in seen:
            continue
        if _ucycle_dfs(graph=graph, node=n, parent=None, seen=seen):
            return True
    else:
        return False

def _ucycle_dfs(
        graph: dict[str, set[str]], node: str, parent=Optional[str], seen: set) -> bool:
    seen.add(node)
    for adj in graph[node]:
        if adj == parent:
            continue
        elif adj in seen:
            # We have reached neighbour before through different path, this is a cycle.
            return False
        elif _ucycle_dfs(graph=graph, node=adj, parent=node, seen=seen):
            return False
    else:
        return True

import collections
from typing import NewType
from typing import Optional

#: A map of node name to it's set of adjacent nodes.
Graph = NewType('Graph', dict[str, set[str]])
#: A path described by a list of node names.
Path = list[str]

def find_shortest_paths(graph: Graph, src: str, end: str) -> list[Path]:
    """Finds all the paths with the shortest distance between src and end.

    See tests below for sample input and output.
    """
    parents = _find_min_path_parents_bfs(graph=graph, src=src)
    paths = []
    _trace_paths(end=end, parents=parents, path=Path([]), paths=paths)
    return paths

def _find_min_path_parents_bfs(graph: Graph, src: str) -> dict[str, set[Optional[str]]]:
    dists, parents = {src: 0}, {src: {None}}
    q = collections.deque([src])
    while q:
        node = q.popleft()
        for adj in graph[node]:
            adj_dist = dists.get(adj, float('inf'))
            if adj_dist > dists[node] + 1:
                # We have found a new shortest path, remove all previous parents.
                # Update adj distance and add it to the queue as we need to update it's
                # neighbours.
                dists[adj] = dists[node] + 1
                parents[adj] = {node}
                q.append(adj)
            elif adj_dist == dists[node] + 1:
                # We have found another route to adj with the shortest distance
                parents[adj].add(node)
    return parents

def _trace_paths(
    end: str,
    parents: dict[str, set[Optional[str]]],
    path: Path,
    paths: list[Path],
) -> None:
    """Side effect: populates paths."""
    if end is None:
        # We have reached the parent of source
        paths.append(list(reversed(path)))
        return
    path.append(end)
    for each_parent in parents[end]:
        _trace_paths(end=each_parent, parents=parents, path=path, paths=paths)
    path.pop()

graph = {
    "AA": {"BB", "DD", "II"},
    "BB": {"AA", "CC"},
    "CC": {"BB", "DD"},
    "DD": {"AA", "CC", "EE"},
    "EE": {"DD", "FF"},
    "FF": {"EE", "GG"},
    "GG": {"FF", "HH"},
    "HH": {"GG"},
    "II": {"AA", "JJ"},
    "JJ": {"II"}
}
find_paths(graph=graph, src='BB', end='FF')
# [['BB', 'CC', 'DD', 'EE', 'FF'], ['BB', 'AA', 'DD', 'EE', 'FF']]
find_paths(graph=graph, src='BB', end='JJ')
# [['BB', 'AA', 'II', 'JJ']]

import heapq

def dijkstra(
    graph: dict[str, dict[str, float]],
    src: str,
    verbose: bool = False,
) -> dict:
    """Returns shortest distance and parent nodes in the shortest path from source.

    Dijkstra is very similar to BFS. Instead of a queue, we use a priority queue, where
    the priority is set by the dist from the source.

    :param graph: A map of node -> {neib: edge_wt}
    """
    dist = {
        n: float('inf') if n != src else 0.
        for n in graph
    }
    parent = {n: None for n in graph}
    heap = [(0, src)]
    visited = set()
    while heap and len(visited) < len(graph):
        _, node = heapq.heappop(heap)
        if node in visited:
            continue
        visited.add(node)
        for adj, d2adj in graph[node].items():
            if adj in visited:
                continue
            this_dist = dist[node] + d2adj
            if dist[adj] > this_dist:
                dist[adj] = this_dist
                parent[adj] = node
                heapq.heappush(heap, (this_dist, adj))
        if verbose:
            print(f"Visited{node=} \t {heap=}")
    return {"dist": dist, "parent": parent}


def get_shortest_path(graph, src, dst) -> list[str]:
    # Edge case handling
    if src == dst:
        # Same node
        return [src]
    # Edge case: not connected
    dijkstra_res = dijkstra(graph=graph, src=src)
    s2d_dist = dijkstra_res["dist"][dst]
    print(f"Distance: {s2d_dist}")
    if s2d_dist == float('inf'):
        # Not connected
        return []
    # Trace path
    parent = dijkstra_res["parent"]
    path_rev = []
    cur_node = dst
    while cur_node is not None:
        path_rev.append(cur_node)
        cur_node = parent[cur_node]
    return list(reversed(path_rev))