Yen算法的python实现

Yen算法是一种用于在加权图中寻找K条最短路径的算法。以下是Yen算法的Python实现：

import heapq

def yen_ksp(graph, start, end, k):
    A = []
    B = []
    A.append(dijkstra(graph, start, end))
    for k_i in range(1, k):
        for i in range(len(A[k_i-1])-1):
            spurNode = A[k_i-1][i]
            rootPath = A[k_i-1][:i+1]
            for path in A:
                if rootPath == path[:i+1]:
                    graph[path[i]][path[i+1]]['weight'] = float('inf')
            for n in range(len(rootPath)-1):
                node = rootPath[n]
                edges = list(graph.edges(node, data=True))
                for edge in edges:
                    if edge[2]['weight'] < 0:
                        raise ValueError("Graph contains negative weight edges")
                for edge in edges:
                    if edge[1] == rootPath[n+1]:
                        graph.remove_edge(node, rootPath[n+1])
                if n == len(rootPath)-2:
                    spurGraph = graph.copy()
                    for node in rootPath[:-1]:
                        for edge in list(spurGraph.edges(node, data=True)):
                            spurGraph.remove_edge(edge[0], edge[1])
                    try:
                        spurPath = dijkstra(spurGraph, spurNode, end)
                        totalPath = rootPath[:-1] + spurPath
                        if totalPath not in B:
                            B.append(totalPath)
                    except:
                        pass
            for path in A:
                if rootPath == path[:i+1]:
                    graph.add_edge(path[i], path[i+1], weight=path[i][path[i+1]]['weight'])
        if len(B) == 0:
            break
        B = sorted(B, key=lambda x: sum(graph[x[i]][x[i+1]]['weight'] for i in range(len(x)-1)))
        A.append(B[0])
        B.pop(0)
    return A

def dijkstra(graph, start, end):
    heap = [(0, start, [])]
    visited = set()
    while heap:
        (cost, node, path) = heapq.heappop(heap)
        if node not in visited:
            visited.add(node)
            path = path + [node]
            if node == end:
                return path
            for next_node in graph.neighbors(node):
                weight = graph[node][next_node]['weight']
                if next_node not in visited:
                    heapq.heappush(heap, (cost+weight, next_node, path))
    return None