从Scipy生成的MST中恢复显式零权重边的方法问询

我有一个包含零权重边的图，需要计算它的最小生成树（MST）。示例图的MST应该包含(0,2)这条零权重边，但用Scipy的minimum_spanning_tree函数计算后，结果里没有这条边。以下是我的代码和输出，请问怎么从输出的MST里恢复这条信息？

输入图与最小生成树示意图 #

input graph             minimum spanning tree

     (0)                         (0)
    / | \                       / | \
   2  |  3                     2  |  3
  /   |   \                   /   |   \
(3)----5--(1)               (3)   |    (1)
  \   |    /                      0     
   2  0   3                       |    
    \ | /                         |  
     (2)                         (2)

代码 #

from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import minimum_spanning_tree
X = csr_matrix([[0, 3, 0, 2],
                [3, 0, 3, 5],
                [0, 3, 0, 2],
                [2, 5, 2, 0]])
X[0,2] = 0
X[2,0] = 0
Tcsr = minimum_spanning_tree(X)
print(Tcsr)
print(f'non-zero entries: {Tcsr.nnz}')

输出 #

(0, 1)  3.0
(0, 3)  2.0
non-zero entries: 2

解答 #

核心原因 #

问题本质是稀疏矩阵的存储特性导致的：Scipy的csr_matrix只存储非零元素，当你直接给csr_matrix的某个位置赋值0时，如果该位置原本没有存储值（初始就是0），这个赋值操作不会被记录——也就是说(0,2)这条零权重边根本没被加入到图的结构里，minimum_spanning_tree自然不会处理它。

另外，即使图正确加载，Prim算法（Scipy实现的MST算法）在存在多个等价MST时，可能会选择其他权重相同的边，但零权重是最小的，这种情况只会发生在零边连接的两个节点已经被其他路径连通的场景。

解决方法 #

1. 正确添加零权重边到稀疏矩阵

改用lil_matrix（支持灵活修改的稀疏矩阵格式）添加零边，再转回csr_matrix：

from scipy.sparse import csr_matrix, lil_matrix
from scipy.sparse.csgraph import minimum_spanning_tree

# 用lil_matrix创建初始图，方便修改
X = lil_matrix([[0, 3, 0, 2],
                [3, 0, 3, 5],
                [0, 3, 0, 2],
                [2, 5, 2, 0]])
# 添加0权重边，lil_matrix会保留这个0值的位置
X[0,2] = 0
X[2,0] = 0
# 转回csr_matrix用于计算MST
X_csr = X.tocsr()
Tcsr = minimum_spanning_tree(X_csr)
print(Tcsr)
print(f'non-zero entries: {Tcsr.nnz}')

执行后输出会包含(0,2)这条零边，边数为3（符合4个节点的MST边数要求）。

2. 手动补充等价MST中的零权重边

如果确认图已正确加载，但算法因等价MST选择了其他边，可以通过连通性检查补充零边：

from scipy.sparse.csgraph import connected_components

# 获取当前MST的连通分量标记
_, component_labels = connected_components(Tcsr)
# 遍历所有零权重边
zero_edges = [(i,j) for i in range(X_csr.shape[0]) 
              for j in range(i+1, X_csr.shape[1]) 
              if X_csr[i,j] == 0]
# 检查零边的两个节点是否在同一连通分量，是则加入MST
for u, v in zero_edges:
    if component_labels[u] == component_labels[v]:
        Tcsr[u, v] = 0
        Tcsr[v, u] = 0
# 清理矩阵中的零元素标记
Tcsr.eliminate_zeros()
print(Tcsr)

这样就能把所有符合条件的零权重边补充到MST中，得到你期望的结果。

内容的提问来源于stack exchange，提问作者FO1234