np.argsort which excludes zero values

问题

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?

回答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])

回答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])

回答3:

You can try:

pd.Series([0.2,0,0,0,0.3,0,0,0,0.4]).sort_values()[lambda x: x != 0].index.values
# array([0, 4, 8])

Or use numpy:

ind = arr.argsort()
ind[arr[ind] != 0]

# array([0, 4, 8])

回答4:

You can cheat a bit with numpy.where, since it defaults to condition.nonzero() anyway. Use numpy.in1d to build a mask.

x=[0.2,0,0,0,0.3,0,0,0,0.4]

np.argsort(x)[np.in1d(np.argsort(x),np.where(x),1)]

Out[35]: array([0, 4, 8], dtype=int32)

来源：https://stackoverflow.com/questions/40857349/np-argsort-which-excludes-zero-values