python - 稀疏矩阵 : how to get nonzero indices for each row

Question

我有一个 scipy CSR 矩阵，我想获取每一行的元素列索引。我的做法是:

import scipy.sparse as sp
N = 100
d = 0.1
M = sp.rand(N, N, d, format='csr')

indM = [row.nonzero()[1] for row in M]

indM 是我需要的，它的行数与 M 相同，看起来像这样:

[array([ 6,  7, 11, ..., 79, 85, 86]),
 array([12, 20, 25, ..., 84, 93, 95]),
...
 array([ 7, 24, 32, 40, 50, 51, 57, 71, 74, 96]),
 array([ 1,  4,  9, ..., 71, 95, 96])]

问题是对于大矩阵，这种方法看起来很慢。有什么方法可以避免列表理解或以某种方式加快它的速度吗？

谢谢。

Answer 1

您可以直接使用 indices 和 indptr 属性:

import numpy
import scipy.sparse

N = 5
d = 0.3
M = scipy.sparse.rand(N, N, d, format='csr')
M.toarray()
# array([[ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ],
#        [ 0.        ,  0.        ,  0.        ,  0.        ,  0.30404632],
#        [ 0.63503713,  0.        ,  0.        ,  0.        ,  0.        ],
#        [ 0.68865311,  0.81492098,  0.        ,  0.        ,  0.        ],
#        [ 0.08984168,  0.87730292,  0.        ,  0.        ,  0.18609702]])

M.indices
# array([1, 2, 4, 3, 0, 1, 4], dtype=int32)
M.indptr
# array([0, 3, 4, 6, 6, 7], dtype=int32)

numpy.split(M.indices, M.indptr)[1:-1]
# [array([], dtype=int32),
#  array([4], dtype=int32),
#  array([0], dtype=int32),
#  array([0, 1], dtype=int32),
#  array([0, 1, 4], dtype=int32)]