NumPy – Captures the index of the first three consecutive negative numbers… here is a solution to the problem.
NumPy – Captures the index of the first three consecutive negative numbers
I need to find the index where three consecutive negative numbers appear for the first time. In normal Python way, I would do this :
a = [1,-1,1,-1,1,-1,1,-1,-1,-1,1,-1,1]
b=0
for i,v in enumerate(a):
if v<0:
b+=1
else:
b=0
if b==3:
break
indx = i-2
Does anyone know how to do this in a smarter NumPy way?
Solution
This is a vectorization solution with the help of convolution –
def first_consecutive_negative_island(a, N=3):
mask = np.convolve(np.less(a,0),np.ones(N,dtype=int))>=N
if mask.any():
return mask.argmax() - N + 1
else:
return None
sample run-
In [168]: a = [1,-1,1,-1,1,-1,1,-1,-1,-1,1,-1,1]
In [169]: first_consecutive_negative_island(a, N=3)
Out[169]: 7
Works wherever the group is-
In [175]: a
Out[175]: [-1, -1, -1, 1, -1, 1, -1, -1, -1, 1, -1, 1]
In [176]: first_consecutive_negative_island(a, N=3)
Out[176]: 0
There are no negative numbers, and it gracefully returns None
–
In [183]: a
Out[183]: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
In [184]: first_consecutive_negative_island(a, N=3)
For exactly three consecutive negative number searches, we can use slices, like this –
def first_consecutive_negative_island_v2(a):
m = np.less(a,0)
mask = m[:-2] & m[1:-1] & m[2:]
if mask.any():
return mask.argmax()
else:
return None
Time –
In [270]: a = np.random.randint(-1,2,(1000000)).tolist()
In [271]: %timeit first_consecutive_negative_island(a, N=3)
10 loops, best of 3: 44.5 ms per loop
In [272]: %timeit first_consecutive_negative_island_v2(a)
10 loops, best of 3: 38.7 ms per loop