用python写一个求解最小斯坦纳树的代码

以下是使用Kruskal算法求解最小斯坦纳树的Python代码：

from collections import defaultdict

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n

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

    def union(self, x, y):
        px, py = self.find(x), self.find(y)
        if px == py:
            return False
        if self.rank[px] < self.rank[py]:
            px, py = py, px
        self.parent[py] = px
        if self.rank[px] == self.rank[py]:
            self.rank[px] += 1
        return True

def min_spanning_tree(graph):
    n = len(graph)
    uf = UnionFind(n)
    edges = []
    for i in range(n):
        for j in range(i + 1, n):
            if graph[i][j] > 0:
                edges.append((i, j, graph[i][j]))
    edges.sort(key=lambda x: x[2])
    tree = defaultdict(list)
    for u, v, w in edges:
        if uf.union(u, v):
            tree[u].append((v, w))
            tree[v].append((u, w))
    return tree

def min_steiner_tree(n, terminals, edges):
    # 取出所有的终端点
    t = len(terminals)
    if t == 1:
        return {terminals[0]: []}
    # 构造子图
    subgraph = [[0] * t for _ in range(t)]
    for i in range(t):
        for j in range(i + 1, t):
            if (terminals[i], terminals[j]) in edges:
                subgraph[i][j] = subgraph[j][i] = edges[(terminals[i], terminals[j])]
    # 求解子图的最小生成树
    subtree = min_spanning_tree(subgraph)
    # 构造最小斯坦纳树
    steiner = defaultdict(list)
    for u, v, w in edges:
        if u in subtree and v in subtree[u]:
            steiner[u].append((v, w))
            steiner[v].append((u, w))
        elif v in subtree and u in subtree[v]:
            steiner[v].append((u, w))
            steiner[u].append((v, w))
        else:
            for i in range(t):
                if (u, terminals[i]) in subtree and (v, terminals[i]) in subtree:
                    steiner[terminals[i]].append((u, w + subtree[(u, terminals[i])]) if (u, terminals[i]) in subtree else (v, w + subtree[(v, terminals[i])]))
                    steiner[u].append((terminals[i], w + subtree[(u, terminals[i])]) if (u, terminals[i]) in subtree else (terminals[i], w + subtree[(v, terminals[i])]))
                    steiner[v].append((terminals[i], w + subtree[(v, terminals[i])]) if (v, terminals[i]) in subtree else (terminals[i], w + subtree[(u, terminals[i])]))
    return steiner

# 示例代码
if __name__ == '__main__':
    n = 5
    terminals = [0, 2, 4]
    edges = {
        (0, 1): 3,
        (0, 2): 1,
        (1, 2): 1,
        (1, 3): 2,
        (2, 3): 8,
        (2, 4): 1,
        (3, 4): 2
    }
    steiner = min_steiner_tree(n, terminals, edges)
    print(steiner)

其中，`min_spanning_tree`函数使用Kruskal算法求解最小生成树，`min_steiner_tree`函数则使用贪心算法构造最小斯坦纳树。