如何用含索引的二维NumPy数组提取稀疏矩阵的非零元素？

问题原因 #

你用A[B]得到不符合预期的形状，是因为NumPy的索引规则导致的：当用二维数组B作为索引时，NumPy会把B的每一行当作行索引集合，取出对应行的整行数据，所以最终得到(83068, 2, 3000)的结果——83068组，每组包含2个完整行（每行3000个元素），这显然不是你要的单个非零元素集合。

针对不同矩阵类型的解决方案 #

1. 如果A是NumPy稠密数组

你需要把B中的行索引和列索引分开，分别作为二维数组的行、列索引输入：

import numpy as np

# 提取所有行索引和列索引
row_indices = B[:, 0]
col_indices = B[:, 1]
# 提取对应位置的元素，得到一维数组
non_zero_vals = A[row_indices, col_indices]
# 转成二维矩阵（比如列向量形式）
C = non_zero_vals.reshape(-1, 1)

此时C的形状为(83068, 1)，完全符合仅包含非零元素的二维矩阵需求。如果需要其他二维形状，可调整reshape的参数（比如reshape(1, -1)得到行向量）。

2. 如果A是SciPy稀疏矩阵（如csr_matrix、csc_matrix）

稀疏矩阵本身提供了直接获取非零元素的属性，不需要通过索引数组B：

from scipy.sparse import csr_matrix

# 直接获取所有非零元素的一维数组
non_zero_vals = A.data
# 转成二维矩阵
C = non_zero_vals.reshape(-1, 1)

这种方式效率更高，因为稀疏矩阵内部已经存储了非零元素的集合，不需要额外的索引查找。

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