Python 图论算法手册

# /*
# Python 图论算法手册
# --------------------------------------------------------
# 说明：
# 本文件汇总常用图论算法与数据结构实现，包含：
# - 图的邻接表表示与辅助工具
# - 遍历：BFS / DFS（递归与迭代）
# - 最短路：Dijkstra / Bellman-Ford / Floyd-Warshall
# - 最小生成树：Prim（堆优化）/ Kruskal（并查集）
# - 拓扑排序：Kahn（入度） / DFS 版
# - 强连通分量：Kosaraju / Tarjan
# - 辅助：并查集、优先队列示例、图连通性检测
#
# 使用方式：
# 将整个文件复制到本地 .py 文件中运行，底部提供示例与测试。
# */

# --------------------------------------------------------
# 基础：加权与无权图的邻接表封装
# --------------------------------------------------------
from collections import deque, defaultdict
import heapq
import random
import sys

class Graph:
    """无向或有向图（邻接表）
    - 若为有向图，add_edge 参数 directed=True
    - 权值默认 None 表示无权图
    """
    def __init__(self, n=0, directed=False):
        self.n = n
        self.directed = directed
        self.adj = [[] for _ in range(n)]

    def add_edge(self, u, v, w=None):
        self.adj[u].append((v, w))
        if not self.directed:
            self.adj[v].append((u, w))

    def edges(self):
        """返回边列表 (u,v,w)，无向图只返回 u<v 的边"""
        res = []
        for u in range(self.n):
            for v, w in self.adj[u]:
                if self.directed or u < v:
                    res.append((u, v, w))
        return res

    def __repr__(self):
        return str(self.adj)

# --------------------------------------------------------
# 遍历：BFS / DFS
# --------------------------------------------------------
def bfs(graph: Graph, start):
    visited = [False] * graph.n
    order = []
    q = deque([start])
    visited[start] = True
    while q:
        u = q.popleft()
        order.append(u)
        for v, _ in graph.adj[u]:
            if not visited[v]:
                visited[v] = True
                q.append(v)
    return order

def dfs_recursive(graph: Graph, start, visited=None, order=None):
    if visited is None:
        visited = [False] * graph.n
    if order is None:
        order = []
    visited[start] = True
    order.append(start)
    for v, _ in graph.adj[start]:
        if not visited[v]:
            dfs_recursive(graph, v, visited, order)
    return order

def dfs_iterative(graph: Graph, start):
    visited = [False] * graph.n
    stack = [start]
    order = []
    while stack:
        u = stack.pop()
        if visited[u]:
            continue
        visited[u] = True
        order.append(u)
        # 保持与递归相近顺序，逆序压栈
        for v, _ in reversed(graph.adj[u]):
            if not visited[v]:
                stack.append(v)
    return order

# --------------------------------------------------------
# 最短路：Dijkstra（非负权）
# --------------------------------------------------------
def dijkstra(graph: Graph, src):
    dist = [float('inf')] * graph.n
    prev = [None] * graph.n
    dist[src] = 0
    pq = [(0, src)]
    while pq:
        d, u = heapq.heappop(pq)
        if d != dist[u]:
            continue
        for v, w in graph.adj[u]:
            if w is None:
                continue
            nd = d + w
            if nd < dist[v]:
                dist[v] = nd
                prev[v] = u
                heapq.heappush(pq, (nd, v))
    return dist, prev

# --------------------------------------------------------
# 最短路：Bellman-Ford（允许负权，检测负环）
# --------------------------------------------------------
def bellman_ford(graph: Graph, src):
    dist = [float('inf')] * graph.n
    dist[src] = 0
    for _ in range(graph.n - 1):
        updated = False
        for u, v, w in graph.edges():
            if w is None:
                continue
            if dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
                updated = True
            if not graph.directed:
                if dist[v] + w < dist[u]:
                    dist[u] = dist[v] + w
                    updated = True
        if not updated:
            break
    # 负环检测（简单版）
    for u, v, w in graph.edges():
        if w is None:
            continue
        if dist[u] + w < dist[v]:
            return None  # 存在负权回路，返回 None
    return dist

# --------------------------------------------------------
# 最短路：Floyd-Warshall（多源最短路）
# --------------------------------------------------------
def floyd_warshall(graph: Graph):
    n = graph.n
    dist = [[float('inf')] * n for _ in range(n)]
    for i in range(n):
        dist[i][i] = 0
    for u in range(n):
        for v, w in graph.adj[u]:
            if w is None:
                continue
            dist[u][v] = min(dist[u][v], w)
    for k in range(n):
        for i in range(n):
            for j in range(n):
                if dist[i][k] + dist[k][j] < dist[i][j]:
                    dist[i][j] = dist[i][k] + dist[k][j]
    return dist

# --------------------------------------------------------
# 并查集（Union-Find）工具
# --------------------------------------------------------
class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0]*n

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, a, b):
        ra = self.find(a); rb = self.find(b)
        if ra == rb: return False
        if self.rank[ra] < self.rank[rb]:
            self.parent[ra] = rb
        elif self.rank[ra] > self.rank[rb]:
            self.parent[rb] = ra
        else:
            self.parent[rb] = ra
            self.rank[ra] += 1
        return True

# --------------------------------------------------------
# 最小生成树：Kruskal（并查集）
# --------------------------------------------------------
def kruskal(graph: Graph):
    edges = graph.edges()
    # 过滤无权边
    weighted_edges = [(u, v, w) for (u, v, w) in edges if w is not None]
    weighted_edges.sort(key=lambda x: x[2])
    uf = UnionFind(graph.n)
    mst = []
    total = 0
    for u, v, w in weighted_edges:
        if uf.union(u, v):
            mst.append((u, v, w))
            total += w
            if len(mst) == graph.n - 1:
                break
    return mst, total

# --------------------------------------------------------
# 最小生成树：Prim（堆优化）
# --------------------------------------------------------
def prim(graph: Graph, start=0):
    visited = [False]*graph.n
    pq = []
    mst = []
    total = 0
    visited[start] = True
    for v, w in graph.adj[start]:
        if w is not None:
            heapq.heappush(pq, (w, start, v))
    while pq and len(mst) < graph.n - 1:
        w, u, v = heapq.heappop(pq)
        if visited[v]:
            continue
        visited[v] = True
        mst.append((u, v, w))
        total += w
        for nxt, w2 in graph.adj[v]:
            if w2 is not None and not visited[nxt]:
                heapq.heappush(pq, (w2, v, nxt))
    return mst, total

# --------------------------------------------------------
# 拓扑排序：Kahn（基于入度）与 DFS 版
# --------------------------------------------------------
def topo_kahn(graph: Graph):
    indeg = [0]*graph.n
    for u in range(graph.n):
        for v, _ in graph.adj[u]:
            indeg[v] += 1
    q = deque([u for u in range(graph.n) if indeg[u]==0])
    order = []
    while q:
        u = q.popleft()
        order.append(u)
        for v, _ in graph.adj[u]:
            indeg[v] -= 1
            if indeg[v] == 0:
                q.append(v)
    if len(order) != graph.n:
        return None  # 存在环
    return order

def topo_dfs(graph: Graph):
    visited = [0]*graph.n  # 0=未,1=进行中,2=完成
    order = []
    cycle = False
    def dfs(u):
        nonlocal cycle
        visited[u] = 1
        for v, _ in graph.adj[u]:
            if visited[v] == 0:
                dfs(v)
            elif visited[v] == 1:
                cycle = True
        visited[u] = 2
        order.append(u)
    for u in range(graph.n):
        if visited[u]==0:
            dfs(u)
    if cycle:
        return None
    return list(reversed(order))

# --------------------------------------------------------
# 强连通分量：Kosaraju 与 Tarjan
# --------------------------------------------------------
def kosaraju_scc(graph: Graph):
    # 1) DFS 得到出序
    visited = [False]*graph.n
    order = []
    def dfs1(u):
        visited[u]=True
        for v, _ in graph.adj[u]:
            if not visited[v]:
                dfs1(v)
        order.append(u)
    for u in range(graph.n):
        if not visited[u]:
            dfs1(u)
    # 2) 构造反向图
    rg = Graph(graph.n, directed=True)
    for u in range(graph.n):
        for v, _ in graph.adj[u]:
            rg.add_edge(v, u, None)
    # 3) 按出序倒序 DFS
    visited = [False]*graph.n
    components = []
    def dfs2(u, comp):
        visited[u]=True
        comp.append(u)
        for v, _ in rg.adj[u]:
            if not visited[v]:
                dfs2(v, comp)
    for u in reversed(order):
        if not visited[u]:
            comp=[]
            dfs2(u, comp)
            components.append(comp)
    return components

def tarjan_scc(graph: Graph):
    index = 0
    indices = [-1]*graph.n
    low = [0]*graph.n
    stack = []
    onstack = [False]*graph.n
    comps = []
    def strong(u):
        nonlocal index
        indices[u] = index
        low[u] = index
        index += 1
        stack.append(u)
        onstack[u]=True
        for v, _ in graph.adj[u]:
            if indices[v] == -1:
                strong(v)
                low[u] = min(low[u], low[v])
            elif onstack[v]:
                low[u] = min(low[u], indices[v])
        if low[u] == indices[u]:
            comp=[]
            while True:
                w = stack.pop()
                onstack[w]=False
                comp.append(w)
                if w==u:
                    break
            comps.append(comp)
    for u in range(graph.n):
        if indices[u]==-1:
            strong(u)
    return comps

# --------------------------------------------------------
# 辅助：图连通性检测
# --------------------------------------------------------
def is_connected(graph: Graph):
    visited = [False]*graph.n
    def dfs(u):
        visited[u]=True
        for v,_ in graph.adj[u]:
            if not visited[v]:
                dfs(v)
    dfs(0)
    return all(visited)

# --------------------------------------------------------
# 示例与测试（随机图）
# --------------------------------------------------------
if __name__ == "__main__":
    # 构造带权无向图示例
    g = Graph(6, directed=False)
    edges = [
        (0,1,7),(0,2,9),(0,5,14),
        (1,2,10),(1,3,15),(2,3,11),
        (2,5,2),(3,4,6),(4,5,9)
    ]
    for u,v,w in edges:
        g.add_edge(u,v,w)

    print("图邻接表：", g.adj)
    print("\nBFS from 0:", bfs(g,0))
    print("DFS recursive from 0:", dfs_recursive(g,0))
    print("DFS iterative from 0:", dfs_iterative(g,0))

    dist, prev = dijkstra(g, 0)
    print("\nDijkstra distances from 0:", dist)

    mst_k, total_k = kruskal(g)
    print("\nKruskal MST edges:", mst_k, "total weight:", total_k)

    mst_p, total_p = prim(g, start=0)
    print("Prim MST edges:", mst_p, "total weight:", total_p)

    print("\nFloyd-Warshall all-pairs shortest path matrix:")
    fw = floyd_warshall(g)
    for row in fw:
        print(row)

    # 有向图示例用于拓扑与 SCC
    dg = Graph(6, directed=True)
    dg.add_edge(5,2); dg.add_edge(5,0); dg.add_edge(4,0)
    dg.add_edge(4,1); dg.add_edge(2,3); dg.add_edge(3,1)
    print("\nDirected Graph adj:", dg.adj)
    print("Topological (Kahn):", topo_kahn(dg))
    print("Topological (DFS):", topo_dfs(dg))

    print("\nKosaraju SCC:", kosaraju_scc(dg))
    print("Tarjan SCC:", tarjan_scc(dg))

# --------------------------------------------------------
# 结语：
# --------------------------------------------------------
# 本手册收录了常用图论算法的 Python 实现，适合教学与工程入门。
# 可在此基础上扩展：最大流（Edmonds-Karp / Dinic）、最小费用最大流、
# 欧拉回路、哈密顿回路检测、图的可视化等专题。
# */