在numpy中对唯一元素的索引进行分组

Question

我有很多很大的整数列表（超过1亿），里面有很多重复的数字。我想找出每个数字出现的位置。目前我正在做的事情是这样的：

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

这是一个简单快捷的解决方案。

a = np.array([0, 0, 0, 1, 1, 3, 3, 3, 2, 2, 2, 0, 0, 1, 4])
sort_idx = np.argsort(a)
unique, counts = np.unique(a, return_counts=True)
b = {key: sort_idx[sum(counts[:key]): sum(counts[:key]) + counts[key]] for key in unique}

Answer 2

numpy_indexed这个包（声明：我是它的作者）提供了一种解决方案，灵感来自Jaime的想法；不过它有测试、友好的界面，还有很多相关的功能：

import numpy_indexed as npi
unique, idx_groups = npi.group_by(a, np.arange(len(a))

Answer 3

这个问题可以通过使用Python的Pandas库来解决，Pandas是一个用于数据分析的工具，然后用到一个叫DataFrame.groupby的功能。

我们来看下面的例子：

 a = np.array([1, 2, 6, 4, 2, 3, 2])

 import pandas as pd
 df = pd.DataFrame({'a':a})

 gg = df.groupby(by=df.a)
 gg.groups

输出结果：

 {1: [0], 2: [1, 4, 6], 3: [5], 4: [3], 6: [2]}

Answer 4

def to_components(index):
    return np.split(np.argsort(index), np.cumsum(np.unique(index, return_counts=True)[1]))

当然可以！请把你想要翻译的内容发给我，我会帮你把它变得更简单易懂。

Answer 5

这段内容和之前在这里提问的内容很相似，所以接下来我会根据我在那里的回答进行调整。最简单的向量化方法是使用排序。以下代码借鉴了即将发布的1.9版本中np.unique的实现，这个版本包含了独特项计数的功能，具体可以查看这里：

>>> 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并没有计算最后一个独特项的出现次数，因为这对分割索引数组并不需要。如果你想要获取所有的值，可以这样做：

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