如何带行号打印coo_matrix稀疏矩阵的行求和结果？

问题说明 #

我有一个scipy.sparse.coo.coo_matrix类型的稀疏矩阵，内容如下：

type(matrix)
# scipy.sparse.coo.coo_matrix

print(matrix)
#   (0, 14)   1.0
#   (1, 17)   1.0
#   (3, 6)    1.0
#   (6, 3)    1.0
#   (12, 15)  2.0
#   (12, 13)  3.0
#   (13, 16)  2.0
#   (13, 12)  3.0
#   (14, 0)   1.0
#   (15, 16)  1.0
#   (15, 12)  2.0
#   (16, 15)  1.0
#   (16, 13)  2.0
#   (17, 1)   1.0
#   (18, 19)  1.0
#   (19, 18)  1.0

按行求和后，输出结果没有行号，无法明确对应关系：

print(matrix.sum(axis=1))
# [[1.]
#  [1.]
#  [0.]
#  [1.]
#  [0.]
#  [0.]
#  [1.]
#  [0.]
#  [0.]
#  [0.]
#  [0.]
#  [0.]
#  [5.]
#  [5.]
#  [1.]
#  [3.]
#  [3.]
#  [1.]
#  [1.]
#  [1.]]

希望输出仅包含有非零和的行，格式如下：

0: 1
1: 1
3: 1
6: 1
12: 5
13: 5
14: 1
15: 3
16: 3
17: 1
18: 1
19: 1

我尝试了以下代码，但结果不符合预期：

Mc=matrix.tocoo()
Mc_sum = {k:v for k,v in zip(Mc.col, Mc.sum(axis = 1))}

输出类似：

{17: matrix([[63.]]), 11: matrix([[25.]]), 5: matrix([[105.]]), 20: matrix([[31.]])...}

解决方案 #

方法一：筛选非零和的行并输出 #

先计算所有行的和，再筛选出和不为0的行，按格式打印：

import numpy as np

# 计算每行的和，转换为一维numpy数组
row_sums = matrix.sum(axis=1).A.flatten()

# 获取所有和不为0的行号
non_zero_rows = np.where(row_sums != 0)[0]

# 按要求格式输出
for row_idx in non_zero_rows:
    print(f"{row_idx}: {int(row_sums[row_idx])}")

• matrix.sum(axis=1)返回形状为(n,1)的稀疏矩阵，.A将其转为numpy数组，.flatten()压缩为一维数组便于遍历
• np.where(row_sums != 0)[0]筛选出所有行和不为0的行索引
• 遍历行索引，打印行号与对应的和（用int()转换是因为示例中均为整数，若需保留浮点数可直接输出row_sums[row_idx]）

方法二：利用COO矩阵的row属性 #

COO矩阵的row属性存储了所有非零元素的行索引，去重后即可得到所有有非零值的行（若行内元素和不为0，可直接用这些行号）：

import numpy as np

# 获取所有非零元素的行号并去重
unique_rows = np.unique(matrix.row)

# 计算每行的和并转为一维数组
row_sums = matrix.sum(axis=1).A.flatten()

# 输出结果
for row_idx in unique_rows:
    print(f"{row_idx}: {int(row_sums[row_idx])}")

这种方法更贴合COO矩阵的特性，无需额外筛选非零和的行（若存在行内元素相加为0的情况，仍需结合方法一的筛选逻辑）

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