CODEWITHSUNDEEP
pythonscipysparse-matrix
swhat
swhat

Reputation: 493

Take the max along an axis of a sparse matrix after already calculating the argmax along that axis

I want to take both the argmax and max along an axis of a scipy.sparse matrix X

>>> type(X)
scipy.sparse.csr.csr_matrix

>>> idx = X.argmax(axis=0)

>>> maxes = X.max(axis=0)

I don't want to have to compute the max twice, but I can't use the same approach to this as if X were a np.ndarray. How can I apply the indices from argmax to X?

Upvotes: 3

Views: 277

Answers (1)

Hemerson Tacon
Hemerson Tacon

Reputation: 2522

I managed to achieve the result that you want adapting the approach that you linked:

from scipy.sparse import csr_matrix

a = [[4, 0, 0], [0, 3, 0], [0, 0, 1]]
a = csr_matrix(a)
idx = a.argmax(axis=0)
m = a.shape[1]
a[idx,np.arange(m)[None,:]].toarray()

Outputs:

array([[4, 3, 1]], dtype=int32)

Upvotes: 1

Related Questions