Python: find the indices of maximum value in each row/column of a compressed sparse matrix

Question

I want to find the index of the largest element in each row and column of a CSR matrix. What function should I use?

Answer 1

The argmax() function of csr_matrix can be used to find the index of the largest element in each row and column of a CSR matrix. You need to use the argument "axis" to specify row and column.

axis=0 means "for each column which row has the max value"

axis=1 means "for each row which column has the max value"

Here is an example:

>>> import numpy as np
>>> from scipy.sparse import csr_matrix
>>> row = np.array([0, 0, 1, 2, 2, 2])
>>> col = np.array([0, 2, 2, 0, 1, 2])
>>> data = np.array([1, 2, 5, 3, 1, 4])
>>> X=csr_matrix((data, (row, col)), shape=(3, 3))
>>> X.toarray()
array([[1, 0, 2],
       [0, 0, 5],
       [3, 1, 4]])
>>> X.argmax(axis=0)
matrix([[2, 2, 1]])
>>> X.argmax(axis=1)
matrix([[2],
        [2],
        [2]])
>>>