Prim 和 Kruskal 算法实现及测试

Prim算法实现及测试：

import heapq

def prim(graph, start):
    visited = [False] * len(graph)
    visited[start] = True
    edges = [(cost, start, to) for to, cost in graph[start]]
    heapq.heapify(edges)
    mst_cost = 0
    mst_edges = []
    while edges:
        cost, frm, to = heapq.heappop(edges)
        if visited[to]: continue
        visited[to] = True
        mst_cost += cost
        mst_edges.append((frm, to))
        for to_next, cost2 in graph[to]:
            if visited[to_next]: continue
            heapq.heappush(edges, (cost2, to, to_next))
    return mst_cost, mst_edges

graph = [[(1, 4), (2, 1)],
         [(0, 4), (2, 2), (3, 5)],
         [(0, 1), (1, 2), (3, 8)],
         [(1, 5), (2, 8)]]

mst_cost, mst_edges = prim(graph, 0)
print(mst_cost) # should print 7
print(mst_edges) # should print [(0, 2), (2, 1), (1, 3)]

Kruskal算法实现及测试：

def kruskal(graph):
    def find(x):
        if parents[x] != x:
            parents[x] = find(parents[x])
        return parents[x]
    edges = [(cost, frm, to) for frm, to, cost in graph]
    edges.sort()
    parents = list(range(len(graph)))
    mst_cost = 0
    mst_edges = []
    for cost, frm, to in edges:
        if find(frm) == find(to): continue
        mst_cost += cost
        mst_edges.append((frm, to))
        parents[find(frm)] = find(to)
    return mst_cost, mst_edges

graph = [(0, 1, 4), (0, 2, 1), (1, 2, 2), (1, 3, 5), (2, 3, 8)]

mst_cost, mst_edges = kruskal(graph)
print(mst_cost) # should print 7
print(mst_edges) # should print [(0, 2), (1, 2), (1, 3)]

测试结果与预期一致。