我知道这是一个老问题，但我最近在做一个类似的事情，其中性能至关重要，所以我对时间进行了广泛的实验。我希望我的发现对社区有益。
Jaime's solution 基于np.unique 是 Python 中最快的算法，但有一个警告：索引没有排序（因为 numpy 使用 quicksort by default）和结果不同于 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
朴素方法、朴素改进和唯一稳定方法的结果是具有排序索引列表的字典。来自唯一快速的索引未排序。
基准代码
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]，因此可以直接比较结果。