S030 · Graphs — Representation, BFS, DFS, Shortest Path

"""graphs.py — the interview canon: BFS, DFS, Dijkstra, topological sort.
 
Run:  python graphs.py
"""
from __future__ import annotations
from collections import defaultdict, deque
import heapq
from typing import Optional
 
 
# ------------------------------------------------------------------
# 1) Adjacency-list graph builder — handles both directed + undirected
# ------------------------------------------------------------------
def build_graph(edges: list[tuple[int, int]], directed: bool = False) -> dict[int, list[int]]:
    g: dict[int, list[int]] = defaultdict(list)
    for u, v in edges:
        g[u].append(v)
        if not directed:
            g[v].append(u)
    return g
 
 
# ------------------------------------------------------------------
# 2) BFS — shortest edge-count path in an unweighted graph
# ------------------------------------------------------------------
def bfs_shortest(graph: dict[int, list[int]], start: int, goal: int) -> Optional[list[int]]:
    """Return the shortest path (as node list) from start to goal, or None."""
    if start == goal:
        return [start]
    seen = {start}
    parent: dict[int, int] = {start: -1}
    q: deque[int] = deque([start])
    while q:
        u = q.popleft()
        for v in graph.get(u, []):
            if v in seen:
                continue
            seen.add(v)
            parent[v] = u
            if v == goal:
                # Reconstruct path
                path = [v]
                while parent[path[-1]] != -1:
                    path.append(parent[path[-1]])
                path.reverse()
                return path
            q.append(v)
    return None
 
 
# ------------------------------------------------------------------
# 3) DFS — recursive and iterative flavors
# ------------------------------------------------------------------
def dfs_recursive(graph: dict[int, list[int]], start: int, seen: set[int] | None = None) -> list[int]:
    seen = seen if seen is not None else set()
    order: list[int] = []
    seen.add(start)
    order.append(start)
    for v in graph.get(start, []):
        if v not in seen:
            order.extend(dfs_recursive(graph, v, seen))
    return order
 
 
def dfs_iterative(graph: dict[int, list[int]], start: int) -> list[int]:
    seen = {start}
    stack = [start]
    order: list[int] = []
    while stack:
        u = stack.pop()
        order.append(u)
        for v in graph.get(u, []):
            if v not in seen:
                seen.add(v)
                stack.append(v)
    return order
 
 
# ------------------------------------------------------------------
# 4) Connected components (undirected graph)
# ------------------------------------------------------------------
def connected_components(graph: dict[int, list[int]], nodes: list[int]) -> list[list[int]]:
    seen: set[int] = set()
    comps: list[list[int]] = []
    for n in nodes:
        if n in seen:
            continue
        comp = dfs_iterative(graph, n)
        seen.update(comp)
        comps.append(comp)
    return comps
 
 
# ------------------------------------------------------------------
# 5) Dijkstra — shortest path with non-negative weights (S029 reuse)
# ------------------------------------------------------------------
def dijkstra(graph: dict[int, list[tuple[int, float]]], start: int) -> dict[int, float]:
    """graph: node -> list of (neighbor, weight). Returns dist[node]."""
    dist: dict[int, float] = {start: 0.0}
    heap: list[tuple[float, int]] = [(0.0, start)]
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist.get(u, float("inf")):
            continue                                            # stale entry
        for v, w in graph.get(u, []):
            nd = d + w
            if nd < dist.get(v, float("inf")):
                dist[v] = nd
                heapq.heappush(heap, (nd, v))
    return dist
 
 
# ------------------------------------------------------------------
# 6) Topological sort (Kahn's algorithm) — dependency order
# ------------------------------------------------------------------
def topological_sort(graph: dict[int, list[int]], nodes: list[int]) -> Optional[list[int]]:
    """Return a valid topological order for the DAG, or None if there's a cycle."""
    indeg: dict[int, int] = {n: 0 for n in nodes}
    for u in nodes:
        for v in graph.get(u, []):
            indeg[v] = indeg.get(v, 0) + 1
    q: deque[int] = deque([n for n in nodes if indeg[n] == 0])
    order: list[int] = []
    while q:
        u = q.popleft()
        order.append(u)
        for v in graph.get(u, []):
            indeg[v] -= 1
            if indeg[v] == 0:
                q.append(v)
    return order if len(order) == len(nodes) else None          # None → cycle detected
 
 
# ------------------------------------------------------------------
# 7) Cycle detection in a directed graph (three-colour DFS)
# ------------------------------------------------------------------
def has_cycle_directed(graph: dict[int, list[int]], nodes: list[int]) -> bool:
    WHITE, GRAY, BLACK = 0, 1, 2
    color = {n: WHITE for n in nodes}
 
    def visit(u: int) -> bool:
        color[u] = GRAY
        for v in graph.get(u, []):
            if color.get(v, WHITE) == GRAY:                     # back-edge → cycle
                return True
            if color.get(v, WHITE) == WHITE and visit(v):
                return True
        color[u] = BLACK
        return False
 
    return any(color[n] == WHITE and visit(n) for n in nodes)
 
 
# ------------------------------------------------------------------
# Runner
# ------------------------------------------------------------------
if __name__ == "__main__":
    # Undirected graph:
    #  1 - 2 - 3
    #  |       |
    #  4 - 5 - 6
    edges = [(1, 2), (2, 3), (1, 4), (4, 5), (5, 6), (3, 6)]
    g = build_graph(edges)
    print("BFS 1 -> 6:", bfs_shortest(g, 1, 6))
    print("DFS from 1 (recursive):", dfs_recursive(g, 1))
    print("DFS from 1 (iterative):", dfs_iterative(g, 1))
    print("connected_components:", connected_components(g, [1, 2, 3, 4, 5, 6, 7, 8]))   # 7 & 8 isolated
 
    # Weighted directed graph for Dijkstra
    weighted = {
        1: [(2, 7.0), (3, 9.0), (6, 14.0)],
        2: [(3, 10.0), (4, 15.0)],
        3: [(4, 11.0), (6, 2.0)],
        4: [(5, 6.0)],
        5: [],
        6: [(5, 9.0)],
    }
    print("Dijkstra from 1:", dijkstra(weighted, 1))
 
    # DAG for topological sort — a tiny ‘build system’
    dag = build_graph([("compile", "link"), ("test", "package"), ("link", "package"), ("bootstrap", "compile"), ("bootstrap", "test")], directed=True)
    all_nodes = ["bootstrap", "compile", "link", "test", "package"]
    print("topological_sort:", topological_sort(dag, all_nodes))
    print("has_cycle (DAG):", has_cycle_directed(dag, all_nodes))
 
    # Add a cycle: package -> bootstrap
    dag2 = build_graph([("compile", "link"), ("test", "package"), ("link", "package"), ("bootstrap", "compile"), ("bootstrap", "test"), ("package", "bootstrap")], directed=True)
    print("has_cycle (with cycle):", has_cycle_directed(dag2, all_nodes))
    print("topological_sort (cyclic):", topological_sort(dag2, all_nodes))