如何筛选CSR格式邻接矩阵中行列属于指定保留节点列表的条目

实现方法 #

你需要筛选行索引或列索引属于保留列表的非零边，直接用X[keep]仅能筛选指定行的内容，无法匹配列符合条件的条目。可以通过将CSR矩阵转换为COO格式快速操作行列索引完成过滤，具体代码如下：

import numpy as np
import scipy.sparse as sp

# 原始邻接矩阵构造
row = np.array([0, 0, 1, 3, 2, 2,3,4,4])
col = np.array([0, 4, 2, 2, 3, 1,4,0,3])
data = np.array([1]*len(row))
X = sp.csr_matrix((data, (row, col)), shape=(5, 5))

# 保留节点列表
keep = [1, 2]

# 转换为COO格式便于直接操作行列索引
X_coo = X.tocoo()
# 构造过滤掩码：行在keep中 或 列在keep中
mask = np.isin(X_coo.row, keep) | np.isin(X_coo.col, keep)
# 用掩码过滤后重建CSR矩阵，保留原始行列尺寸
filtered_X = sp.csr_matrix(
    (X_coo.data[mask], (X_coo.row[mask], X_coo.col[mask])),
    shape=X.shape
)

# 验证输出
print(filtered_X)

运行后输出结果和预期一致：

(1, 2)	1
  (2, 1)	1
  (2, 3)	1
  (3, 2)	1

如果需要将保留的节点重新映射为连续的小索引（比如把1→0、2→1、3→2），额外对过滤后的行列索引做映射处理即可。

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