python - numpy中唯一元素的分组索引

Question

我有许多包含许多重复项的大型(>100,000,000)整数列表。我想获取每个元素出现的索引。目前我正在做这样的事情:

import numpy as np
from collections import defaultdict

a = np.array([1, 2, 6, 4, 2, 3, 2])
d=defaultdict(list)
for i,e in enumerate(a):
    d[e].append(i)

d
defaultdict(<type 'list'>, {1: [0], 2: [1, 4, 6], 3: [5], 4: [3], 6: [2]})

这种遍历每个元素的方法非常耗时。有没有一种有效的或矢量化的方法来做到这一点？

编辑1我在下面尝试了 Acorbe 和 Jaime 的方法

a = np.random.randint(2000, size=10000000)

结果是

original: 5.01767015457 secs
Acorbe: 6.11163902283 secs
Jaime: 3.79637312889 secs

Answer 1

这与问题非常相似here ，所以接下来是我在那里的回答的改编。对其进行矢量化的最简单方法是使用排序。以下代码大量借鉴了 np.unique 的实现，用于即将发布的 1.9 版本，其中包括独特的项目计数功能，请参阅 here :

>>> a = np.array([1, 2, 6, 4, 2, 3, 2])
>>> sort_idx = np.argsort(a)
>>> a_sorted = a[idx]
>>> unq_first = np.concatenate(([True], a_sorted[1:] != a_sorted[:-1]))
>>> unq_items = a_sorted[unq_first]
>>> unq_count = np.diff(np.nonzero(unq_first)[0])

现在:

>>> unq_items
array([1, 2, 3, 4, 6])
>>> unq_count
array([1, 3, 1, 1, 1], dtype=int64)

要获取每个值的位置索引，我们只需执行以下操作:

>>> unq_idx = np.split(sort_idx, np.cumsum(unq_count))
>>> unq_idx
[array([0], dtype=int64), array([1, 4, 6], dtype=int64), array([5], dtype=int64),
 array([3], dtype=int64), array([2], dtype=int64)]

现在您可以构建压缩unq_items 和unq_idx 的字典。

请注意，unq_count 不计算最后一个唯一项的出现次数，因为拆分索引数组不需要它。如果你想拥有所有你可以做的值(value)观:

>>> unq_count = np.diff(np.concatenate(np.nonzero(unq_first) + ([a.size],)))
>>> unq_idx = np.split(sort_idx, np.cumsum(unq_count[:-1]))