Question

我有许多大型（＆gt; 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]})

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

EDIT1 我在下面的

中尝试了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并不计算最后一个唯一项的出现次数，因为拆分索引数组不需要这样做。如果你想拥有所有可以做的值：

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

Answer 2

这可以通过python pandas（python数据分析库）和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 3

index.js包（免责声明：我是它的作者）实现了一个受Jaime's启发的解决方案;但是有了测试，一个很好的界面，以及许多相关的功能：

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

Answer 4

简单快捷的解决方案。

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 5

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

Answer 6

我知道这是一个老问题，但我最近正在研究一个类似的事情，其中性能至关重要，因此我对时间进行了广泛的试验。我希望我的发现对社区有益。

基于 np.unique 的

Jaime's solution 是 Python 中最快的算法，但有一个警告：索引未排序（因为 numpy 使用 quicksort { {3}}) 并且结果不同于 OP 的原始算法（以下称为 naive）。使用 stable 选项可以修复它，但会稍微减慢速度。

可以使用 Python 的内置 array 模块改进朴素的方法，如下所示：

import array
from collections import defaultdict

a = np.array(...)  # 1D, int array
d = defaultdict(lambda: array.array("L"))
alist = array.array("L")
alist.frombytes(a.tobytes())
for n in range(len(alist)):
    d[alist[n]].append(n)

它只是比 Jaime 的稳定排序解决方案慢几分之一。

这是在我的平台上使用 Python 3 完成的一些测试

Best of 5
Naive method: 0.21274029999999988 s
Naive improved: 0.13265090000000002 s
Unique quick: 0.073496 s
Unique stable: 0.1235801999999997 s

naive 方法、naive 改进和唯一稳定的结果都将是具有排序索引列表的字典。独特的快速不会。

基准代码

import array
import timeit
from collections import defaultdict

import numpy as np

def count_naive(a):
    d = defaultdict(list)
    for n, e in enumerate(a):
        d[e].append(n)
    return dict(d)

def count_improved(a):
    d = defaultdict(lambda: array.array("L"))
    alist = array.array("L")
    alist.frombytes(a.tobytes())
    for n in range(len(alist)):
        d[alist[n]].append(n)
    return {n: indices.tolist() for n, indices in d.items()}

def count_unique(a):
    sorted_idx = np.argsort(a)  # , kind='stable')
    counts = np.bincount(a)
    split_idx = np.split(sorted_idx, np.cumsum(counts[:-1]))
    return {n: indices.tolist() for n, indices in enumerate(split_idx)}

def count_stable(a):
    sorted_idx = np.argsort(a, kind="stable")
    counts = np.bincount(a)
    split_idx = np.split(sorted_idx, np.cumsum(counts[:-1]))
    return {n: indices.tolist() for n, indices in enumerate(split_idx)}

a = np.random.randint(1000, size=1000000)

trials = 5
t_naive = timeit.repeat("count_naive(a)", globals=globals(), repeat=trials, number=1)
t_improved = timeit.repeat("count_improved(a)", globals=globals(), repeat=trials, number=1)
t_unique = timeit.repeat("count_unique(a)", globals=globals(), repeat=trials, number=1)
t_stable = timeit.repeat("count_stable(a)", globals=globals(), repeat=trials, number=1)

print(f"Best of {trials}")
print(f"Naive method: {min(t_naive)} s")
print(f"Naive improved: {min(t_improved)} s")
print(f"Unique quick: {min(t_unique)} s")
print(f"Unique stable: {min(t_stable)} s")

注意所有函数都以返回 Dict[int, list] 的方式编写，因此可以直接比较结果。

在numpy中对独特元素的索引进行分组

6 个答案: