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

Question

目标：我有一个大的1d阵列（3000000+）距离，有许多重复距离。 我正在尝试编写最快的函数，它返回在数组中出现n次的所有距离。 我已经在numpy中编写了一个函数，但代码中的一行有一个瓶颈。 Swift性能是一个问题，因为计算是在2400个不同的大距离阵列的for循环中完成的。

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

预期结果：给定1d距离数组[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倍的改进。 要获得更多，您必须在GPU上找到并行解决方案。

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

Answer 2

你可以使用一个Counter ：

>>> a = np.array([2000000000,3005670000,2000000000,12345667,4000789000,12345687,12345667,2000000000,12345667])
>>> c = Counter(a)
>>> np.array([i for i in c if c[i] >= 3])
array([2000000000,   12345667])