如何在數組中找到N個連續數字相等？

Question

我有一個 numpy.ndarray，其值為 -1 或 1。該數組很長，但我會發布一些示例。

array([ 1, -1,  1, -1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1, -1,  1, -1,
    1,  1,  1], dtype=int8)

數組一直持續到索引 200 左右。我想找到有 N 個連續值等於 1 或 -1 的位置。 因此，在上面的示例中，如果 N 等於 8，它應該返回索引 11、12、13。

Answer 1

您可以使用.cumsum來實現滑動 window：

a = np.array([ 1, -1,  1, -1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1, -1,  1, -1,
    1,  1,  1], dtype=np.int8) 
N = 8

A = a.cumsum()
A[N:] -= A[:-N]
(np.abs(A)==N).nonzero()[0]
# array([11, 12, 13])

Answer 2

我想使用列表分享我的幼稚解決方案

讓arr = [1,1,-1,-1,...]

index = []
myList = list(arr)
for i in range(m-n): # m is length of array
    if myList[i:i+n].count(1)==n or myList[i:i+n].count(-1)==n:
        index.append(i)

index[]現在包含滿足您條件的所有索引

Answer 3

這是一個np.convolve -

np.flatnonzero(np.convolve(np.diff(a)==0, [1]*(N-1))==(N-1))+1

樣品運行 -

In [146]: N = 8

In [147]: np.flatnonzero(np.convolve(np.diff(a)==0, [1]*(N-1))==(N-1))+1
Out[147]: array([11, 12, 13])

Answer 4

由於您標記pandas ：

s = pd.Series(a)

N = 8
np.where(s.diff().eq(0).rolling(N-1).sum()==N-1)

Output：

(array([11, 12, 13]),)