CODEWITHSUNDEEP
pythonarraysnumpy
WesH
WesH

Reputation: 470

Numpy: Is there a way to use np.argmax to search for consecutive values?

Is there a way to search a Numpy array for consecutive values and return their index?

For example, if we use argmax on an array to find the first instance of a single element:

import numpy as np

a = np.array((0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1))
print(np.argmax(a == 1))

We get the result:

3

Is there a way to search this same array for the first instance of a pair? E.g. I want to get the value 6 for the following code:

import numpy as np

a = np.array((0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1))
print(np.argmax(a == [1, 1]))

But instead it returns:

0

Upvotes: 1

Views: 179

Answers (1)

Divakar
Divakar

Reputation: 221554

For small number of consecutive search pattern, we can simply slice and look for all matches and finally use our-favourite argmax.

Hence, for 2 consecutive search pattern, it would be -

In [6]: ((a[:-1]==1) & (a[1:]==1)).argmax()
Out[6]: 6

Some explanation

It's all about slicing, as we get two one-offsetted array slices. The one-offset comes because of the window-length of consecutive search being 2. So, for a search window-length of 3, we need to consider two-offsetted array slices and so on. Now, back to our simplistic two-consecutive window case, we have one-offsetted slices. We compare these against 1, which gives us boolean arrays of matches. Then comes AND-ing, so that the entire window is covered. Finally jumps in our argmax to get the start of the first of those islands!

A break-down of the steps should help clarify further on the given sample -

In [24]: a
Out[24]: array([0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1])

In [25]: a[:-1]
Out[25]: array([0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1])

In [26]: a[:-1]==1
Out[26]: 
array([False, False, False,  True, False, False,  True,  True, False,
        True,  True])

In [27]: a[1:]==1
Out[27]: 
array([False, False,  True, False, False,  True,  True, False,  True,
        True,  True])

In [28]: (a[:-1]==1) & (a[1:]==1)
Out[28]: 
array([False, False, False, False, False, False,  True, False, False,
        True,  True])

In [29]: ((a[:-1]==1) & (a[1:]==1)).argmax()
Out[29]: 6

Larger number of consecutive occurrences

For larger number of cons. occurrences, we can resort to something more built-in and use np.convolve, like so -

In [20]: W = 2 # window-length

In [21]: np.convolve(a,[1]*W,'same').argmax()-W//2
Out[21]: 6

In [22]: W = 3

In [23]: np.convolve(a,[1]*W,'same').argmax()-W//2
Out[23]: 9

Upvotes: 1

Related Questions