Fast nonzero indices per row/column for (sparse) 2D numpy array

I am looking for the fastest way to obtain a list of the nonzero indices of a 2D array per row and per column. The following is a working piece of code:

preds = [matrix[:,v].nonzero()[0] for v in range(matrix.shape[1])]
descs = [matrix[v].nonzero()[0] for v in range(matrix.shape[0])]


Example input:

matrix = np.array([[0,0,0,0],[1,0,0,0],[1,1,0,0],[1,1,1,0]])


Example output

preds = [array([1, 2, 3]), array([2, 3]), array([3]), array([], dtype=int64)]
descs = [array([], dtype=int64), array([0]), array([0, 1]), array([0, 1, 2])]


(The lists are called preds and descs because they refer to the predecessors and descendants in a DAG when the matrix is interpreted as an adjacency matrix but this is not essential to the question.)

I was wondering whether this might be doable with some sort of sparse matrix (CSR, CSC, COO etc.) from scipy.sparse but I am unfamiliar with them and have not got that working. I don't necessarily need to use these types if a faster option exists.

Timing example: For timing purposes, the following matrix is a good representative:

test_matrix = np.zeros(shape=(4096,4096),dtype=np.float32)
for k in range(16):
test_matrix[256*(k+1):256*(k+2),256*k:256*(k+1)]=1


Thank you.

Background: In my code, these two lines take 75% of the time for a 4000x4000 matrix whereas the ensuing topological sort and DP algorithm take only the rest of the quarter. If someone know how to do this much more efficiently it would be greatly appreciated. Roughly 5% of the matrix has nonzero values.

(on suggestion I moved the question here from: https://stackoverflow.com/questions/62065793/fast-nonzero-indices-per-row-column-for-sparse-2d-numpy-array Contains several useful answers)

• For quite a few algorithms on large sparse graphs, see scipy.sparse.csgraph. To get rows and columns (not just indices), see scipy.sparse.csr_matrix getrow() and getcol(). – denis Aug 10 '20 at 10:36

1 Answer

Ideally you would have the matrix already in a sparse matrix data structure. But for this example we can do the conversion via

In [31]: M = scipy.sparse.coo_matrix(np.array([[0,0,0,0],[1,0,0,0],[1,1,0,0],[1,1,1,0]]))


Then you can do

In [32]: Mcsr = M.tocsr()

In [33]: np.split(Mcsr.indices, Mcsr.indptr)
Out[33]:
[array([], dtype=int32),
array([], dtype=int32),
array([0], dtype=int32),
array([0, 1], dtype=int32),
array([0, 1, 2], dtype=int32),
array([], dtype=int32)]


to get descs. Similarly,

In [34]: Mcsc = M.tocsc()

In [35]: np.split(Mcsc.indices, Mcsc.indptr)
Out[35]:
[array([], dtype=int32),
array([1, 2, 3], dtype=int32),
array([2, 3], dtype=int32),
array([3], dtype=int32),
array([], dtype=int32),
array([], dtype=int32)]


gives preds. I have no idea if this is more effective in practice though.

• Thank you for your reply! The code does what it is supposed to do and is faster (my earlier comment had a wrong timing). If I run some tests on the test_matrix I added I get 2030 ms (for 10 runs) for my original code and 1170 ms for this code so this is a speed up of 2 for this matrix. – Richard Schoonhoven May 28 '20 at 16:08