python - 在 numpy 中执行间接排序，为相同元素保留相同索引

Question

假设我有一个像这样的数组:

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

我想将它转换为这样的数组:

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

我尝试过 np.argsort，但返回

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

这并不完全符合我需要它做的事情(它仍然为相同的元素分配不同的索引。

到目前为止，我已经编写了这个小函数来完成我想要的事情:

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

是否有一种优雅的方式来使用 numpy/scipy 执行我想要的操作？

Answer 1

这就是 numpy.unique 的 return_inverse 参数。是:

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