返回scipy稀疏矩阵的对角元素

Question

给定scipy.sparse.csr.csr_matrix，是否可以快速返回对角线上的元素？

我想这样做的原因是计算inv(D) A，其中D是对角矩阵，其对角线条目与A一致（A是我的稀疏矩阵，保证对角线上有非零值。

Answer 1

使用csr_matrix.diagonal()：

  

返回矩阵的主对角线

示例：

>>> import numpy as np
>>> from scipy.sparse import csr_matrix
>>> mymat = csr_matrix((4, 4), dtype=np.int8)
>>> mymat.diagonal()
array([0, 0, 0, 0], dtype=int8)