python sparse matrix get maximum values and index

Question

I have a sparse matrix A(equal to 10 * 3 in dense), such as:

print type(A)
<class scipy.sparse.csr.csr_matrix>

print A
(0, 0)  0.0160478743808
(0, 2)  0.0317314165078
(1, 2)  0.0156596521648
(1, 0)  0.0575683686558
(2, 2)  0.0107481166871
(3, 0)  0.0150580924929
(3, 2)  0.0297743235876
(4, 0)  0.0161931803955
(4, 2)  0.0320187296788
(5, 2)  0.0106034409766
(5, 0)  0.0128109177074
(6, 2)  0.0105766993238
(6, 0)  0.0127786088452
(7, 2)  0.00926522256063
(7, 0)  0.0111941023699

The max values for each column is:

print A.max(axis=0)
(0, 0)  0.0575683686558
(0, 2)  0.0320187296788

I would like to get the index corresponding to the column value. I know that the

A.getcol(i).tolist()

will return me a list of each column which allow me to use argmax() function, but this way is really slow. I am wondering is there any descent way to do?

Answer 1

The more efficient way to get the max and argmax values in each matrix column is simply using scipy.sparse native functions:

• max value of A in each matrix columns:

  max_values = A.max(axis=0)

• max arg of A in each matrix column:

  max_args = A.argmax(axis=0)

The same to compute max values and arg max in each matrix row (using axis=1) or to compute max values and arg max of all the matrix (using axis=None).