Python/Numpy 我已经为大型数组编写了最快的代码吗？

Question

目标:我有一个很大的一维数组(3000000+)，其中有许多重复的距离。我正在尝试编写最快的函数，该函数返回数组中出现 n 次的所有距离。我在 numpy 中编写了一个函数，但代码中的一行存在瓶颈。 Swift 性能是一个问题，因为计算是在 for 循环中针对 2400 个不同的大距离数组完成的。

import numpy as np
for t in range(0, 2400):
 a=np.random.randint(1000000000, 5000000000, 3000000)
 b=np.bincount(a,minlength=np.size(a))
 c=np.where(b == 3)[0] #SLOW STATEMENT/BOTTLENECK
 return c

预期结果:给定一维距离数组 [2000000000,3005670000,2000000000,12345667,4000789000,12345687,12345667,2000000000,12345667]我希望在查询时返回 [2000000000,12345667] 的数组以返回在主数组中出现 3 次的所有距离的数组。

我该怎么办？

Answer 1

使用np.unique:

a=np.random.randint(0,10**6,3*10**6)
uniques,counts=np.unique(a,return_counts=True)
In [293]: uniques[counts==14]
Out[293]: array([  4541, 250510, 622471, 665409, 881697, 920586])

这需要不到一秒钟的时间。但我不明白为什么你的 where 语句很慢。对我来说，您的解决方案更快:

In [313]: %timeit b=np.bincount(a,minlength=a.size)
61.5 ms ± 4.82 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [314]: %timeit np.where(b==3)[0]
11.8 ms ± 271 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [315]: %timeit uniques,counts=np.unique(a,return_counts=True)
424 ms ± 6.82 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [316]: %timeit Counter(a)
1.41 s ± 18 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

编辑

@numba.jit()
def count(a,n):
    counters=np.zeros(10**6,np.int32)
    for i in a:
        counters[i] += 1
    res=np.empty_like(counters)
    k = 0    
    for i,j in enumerate(counters):
        if j == n:
            res[k] = i
            k += 1
    return res[:k]        

这个 numba 函数可以给你带来 3 倍的提升。如需更多信息，您必须找到并行解决方案，on GPU for example .

In [26]: %timeit  count(a,13)
23.6 ms ± 1.56 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)