How to find N consecutive numbers are equal in array?

user_13838655 0
Viewed 266

I have a numpy.ndarray that has values -1 or 1. The array is quite long, but I will post a little sample.

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

The array continues till about index 200. I would like to find the location where there are N consecutive values equal to either 1 or -1. So in the example above it should return index 11, 12, 13 if N was equal to 8.

pandas numpy python
4 Answers

You can use .cumsum to implement a sliding 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])
user_72073920

I would like to share my naive solution using lists

let 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[] now holds all indices which satisfy your condition

Admin1

Here's one with np.convolve -

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

Sample run -

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])
user_32938810

Since you tagged pandas:

s = pd.Series(a)

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

Output:

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