Nikolay Ambartsumov Nikolay Ambartsumov - 6 months ago 11
Python Question

Performing an indirect sort in numpy, preserving same indices for identical elements

Let's assume I have an array like this:

a = np.array([5, 2, 13, 13, 222])


I want to convert it to an array like this:

b = np.array([1, 0, 2, 2, 3])


I've tried np.argsort, but returns

np.argsort(np.array([5, 2, 13, 13, 222])) # = array([1, 0, 2, 3, 4])


which doesn't do exactly what I need it to do (it still assigns different indexes to identical elements.

So far I've written this little function to do what I want:

def indexate_array(v):
v_unique = np.unique(v)
result = shape_like(v)
dic = {value: result for value, result in zip(v_unique, np.argsort(v_unique))}
for i, val in enumerate(v):
result[i] = dic[val]
return result


Is there an elegant way to perform the operation I want using numpy/scipy?

Answer

This is what the return_inverse parameter to numpy.unique does:

In [5]: np.unique(a, return_inverse=True)
Out[5]: (array([  2,   5,  13, 222]), array([1, 0, 2, 2, 3]))