python - 在 NumPy 中查找特定长度的连续重复

Question

假设我有一个 NumPy 数组:

a = np.array([0, 1, 2, 2, 3, 4, 5, 5, 6, 7, 8, 9, 9, 9, 10, 11, 12, 13, 13, 13, 14, 15])

我有一个用户指定的长度 m = 2 ，以便查看时间序列中是否有任何重复的长度。在这种情况下，长度为 m = 2 的重复是:

[2, 2]
[5, 5]
[9, 9]
[9, 9]
[13, 13]

并且用户可以将其更改为m = 3 并且长度为m = 3 的重复是:

[9, 9, 9]
[13, 13, 13]

我需要一个函数，它要么返回找到重复的位置的索引，要么返回 None。因此，对于 m = 3，该函数将返回以下起始索引的 NumPy 数组:

[11, 17]

对于 m = 4，该函数将返回 None。什么是最干净、最快的方法来完成这个？

更新请注意，数组不必排序，我们不对排序后的结果感兴趣。我们只想要未排序数组的结果。 m = 2 的结果对于这个数组应该是相同的:

b = np.array([0, 11, 2, 2, 3, 40, 5, 5, 16, 7, 80, 9, 9, 9, 1, 11, 12, 13, 13, 13, 4, 5])

Answer 1

方法 #1

我们可以利用 1D convolution对于矢量化解决方案 -

def consec_repeat_starts(a, n):
    N = n-1
    m = a[:-1]==a[1:]
    return np.flatnonzero(np.convolve(m,np.ones(N, dtype=int))==N)-N+1

样本运行-

In [286]: a
Out[286]: 
array([ 0,  1,  2,  2,  3,  4,  5,  5,  6,  7,  8,  9,  9,  9, 10, 11, 12,
       13, 13, 13, 14, 15])

In [287]: consec_repeat_starts(a, 2)
Out[287]: array([ 2,  6, 11, 12, 17, 18])

In [288]: consec_repeat_starts(a, 3)
Out[288]: array([11, 17])

In [289]: consec_repeat_starts(a, 4)
Out[289]: array([], dtype=int64)

方法 #2

我们还可以利用 binary-erosion -

from scipy.ndimage.morphology import binary_erosion

def consec_repeat_starts_v2(a, n):
    N = n-1
    m = a[:-1]==a[1:]
    return np.flatnonzero(binary_erosion(m,[1]*N))-(N//2)