在numpy数组中按顺序查找相同整数的更快方法

Question

现在，我只是循环使用np.nditer()并与前面的元素进行比较。是否有一种更快的(矢量化)方法？

额外的好处是，我并不总是要到数组的末尾；一旦找到了一个max_len序列，我就完成搜索。

import numpy as np

max_len = 3
streak = 0
prev = np.nan

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

for c in np.nditer(a):
  if c == prev:
      streak += 1
      if streak == max_len:
          print(c)
          break
  else:
      prev = c
      streak = 1

我想过的另一种方法是使用np.diff()，但这只是转移了问题；我们现在正在寻找其结果中的一个零序列。而且，我怀疑它是否会更快，因为它将不得不计算每个整数的差值，而在实践中，序列会在到达列表末尾之前发生。

Answer 1

我开发了一个工作正常的numpy-only版本，但是经过测试后，我发现它的性能很差，因为它不能利用短路。既然这就是你想要的，我在下面描述。但是，有一种更好的方法使用numba，并对代码进行轻微修改。(请注意，所有这些都返回a中的第一个匹配的索引，而不是值本身。我觉得这种方法更灵活。)

@numba.jit(nopython=True)
def find_reps_numba(a, max_len):
    streak = 1
    val = a[0]
    for i in range(1, len(a)):
        if a[i] == val:
            streak += 1
            if streak >= max_len:
                return i - max_len + 1
        else:
            streak = 1
            val = a[i]
    return -1

这比纯Python版本快100倍。

numpy版本使用滚动窗口技巧和argmax戏法。但是，这也比纯Python版本慢了很多，甚至比纯Python版本慢了30倍。

def rolling_window(a, window):
    a = numpy.ascontiguousarray(a)  # This approach requires a C-ordered array
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return numpy.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)

def find_reps_numpy(a, max_len):
    windows = rolling_window(a, max_len)
    return (windows == windows[:, 0:1]).sum(axis=1).argmax()

我在第一个函数的非both版本中测试了这两种功能。(我使用了木星的%%timeit特性进行测试。)

a = numpy.random.randint(0, 100, 1000000)

%%timeit
find_reps_numpy(a, 3)
28.6 ms ± 553 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
find_reps_orig(a, 3)
4.04 ms ± 40.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
find_reps_numba(a, 3)
8.29 µs ± 89.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

请注意，根据函数在a中搜索的深度，这些数字可能会有很大的差异。为了更好地估计预期性能，我们可以每次重新生成一组新的随机数，但是如果不将该步骤包括在计时中，则很难做到这一点。因此，为了在这里进行比较，我包括了生成随机数组所需的时间，而无需运行任何其他操作：

a = numpy.random.randint(0, 100, 1000000)
9.91 ms ± 129 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

a = numpy.random.randint(0, 100, 1000000)
find_reps_numpy(a, 3)
38.2 ms ± 453 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

a = numpy.random.randint(0, 100, 1000000)
find_reps_orig(a, 3)
13.7 ms ± 404 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

a = numpy.random.randint(0, 100, 1000000)
find_reps_numba(a, 3)
9.87 ms ± 124 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

正如您所看到的，find_reps_numba是如此之快，以至于运行numpy.random.randint(0, 100, 1000000)所需时间的差异要大得多--因此在第一次测试和最后一次测试之间出现了虚幻的加速。

因此，这个故事的主要寓意是，numpy解决方案并不总是最好的。有时候，即使是纯Python也更快。在这些情况下，numba在nopython模式下可能是目前为止最好的选择。

Answer 2

您可以从groupby包中使用itertools。

import numpy as np
from itertools import groupby

max_len = 3
best = ()

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

for k, g in groupby(a):
    tup_g = tuple(g)
    if tup_g==max_len:
        best = tup_g
        break
    if len(tup_g) > len(best):
        best = tup_g

best
# returns:
(2, 2, 2)

Answer 3

假设您正在寻找至少连续出现max_len次数的元素，下面是一种基于NumPy的方法-

m = np.r_[True,a[:-1]!=a[1:],True]
idx0 = np.flatnonzero(m)
m2 = np.diff(idx0)>=max_len
out = None # None for no such streak found case
if m2.any():
    out = a[idx0[m2.argmax()]]

另一个用binary-dilation -

from scipy.ndimage.morphology import binary_erosion

m = np.r_[False,a[:-1]==a[1:]]
m2 = binary_erosion(m, np.ones(max_len-1, dtype=bool))
out = None
if m2.any():
    out = a[m2.argmax()]

最后，为了完整起见，您还可以查看numba。您的现有代码将按原样工作，在a上直接循环，即for c in a:。