python - 如何有效地获取唯一值的索引列表？

Question

是否有一种内置方法可以帮助我有效地实现以下目标:给定一个数组，我需要一个数组列表，每个数组都有指向数组的不同唯一值的索引？

如果 f 是想要的函数，

b = f(a)

和

u, idxs = unique(a)

然后

b[i] == where(idxs==i)[0]

我知道 pandas.Series.groupby() 可以做到这一点，但是当有超过 10^5 个唯一整数时创建字典可能效率不高。

Answer 1

如果你有 numpy >= 1.9 你可以这样做:

>>> a = np.random.randint(5, size=10)
>>> a
array([0, 2, 4, 4, 2, 4, 4, 3, 2, 1])
>>> unq, unq_inv, unq_cnt = np.unique(a, return_inverse=True, return_counts=True)
>>> np.split(np.argsort(unq_inv), np.cumsum(unq_cnt[:-1]))
[array([0]), array([9]), array([1, 4, 8]), array([7]), array([2, 3, 5, 6])]
>>> unq
array([0, 1, 2, 3, 4])

在较早的版本中，您可以获得额外的计数:

>>> unq_cnt = np.bincount(unq_inv)

此外，如果您想确保对每个值的索引进行排序，我认为您需要使用稳定排序，例如np.argsort(unq_inv, kind='mergesort')

---

想想你似乎在追求什么，我认为这是最大限度地减少对昂贵函数的调用，我认为你不需要做你要求的事情。假设您的函数是平方函数，您可以简单地执行以下操作:

>>> unq, unq_inv = np.unique(a, return_inverse=True)
>>> f_unq = unq**2
>>> f_a = f_unq[unq_inv]
>>> a
array([0, 2, 4, 4, 2, 4, 4, 3, 2, 1])
>>> f_a
array([ 0,  4, 16, 16,  4, 16, 16,  9,  4,  1])