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() for v in range(matrix.shape)] descs = [matrix[v].nonzero() for v in range(matrix.shape)]
matrix = np.array([[0,0,0,0],[1,0,0,0],[1,1,0,0],[1,1,1,0]])
preds = [array([1, 2, 3]), array([2, 3]), array(), array(, dtype=int64)] descs = [array(, dtype=int64), array(), 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
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)