np.argsort which excludes zero values

Question

I have an array [0.2,0,0,0,0.3,0,0,0,0.4]. I'm using np.argsort to sort values and get that indexes.

So, for my example, it will be something like [1,5,9,2,3,4,6...]. However, I would like to get array of indexes only for non zero values.

In my example only [1,5,9].

How do I implement it in python with pandas and numpy?

Answer 1

Using np.nonzero and indexing trick

def sparse_argsort(arr):
    indices = np.nonzero(arr)[0]
    return indices[np.argsort(arr[indices])]

sparse_argsort(a)

array([0, 4, 8])

one liner:

(lambda a: (lambda a_, i_: i_[np.argsort(a_[i_])])(a,np.nonzero(a)[0]))(a)

array([0, 4, 8])

Answer 2

one line numpy

np.where(a != 0, a, np.nan).argsort()[:(a != 0).sum()]

same logic, two lines, more efficient

nz = a != 0
np.where(nz, a, np.nan).argsort()[:nz.sum()]

array([0, 4, 8])