Find the index of the first occurrence of some value that is not X or Y in a numpy array

user_5861691 0
Viewed 641

I have a vector of the following form:

import numpy as np
vec = np.array([2, 2, 2, 51, 51, 52, 52, 14, 14, 14, 51, 51, 52, 52])

Is there a numpy-thonic way to find the index of the first occurrence of a value that is not (for instance) 51 or 52? In other words, a function that would return the following indexes: [0, 7], where 0 is the index of the first apparition of 2, and 7 is the index of the first apparition of 14.

numpy python
2 Answers

np.unique returns the first index of each number if you specify return_index=True. You can filter the result pretty easily using, e.g., np.isin:

u, i =  np.unique(vec, return_index=True)
result = i[np.isin(u, [51, 52], invert=True)]

The advantage of doing it this way is that u is a significantly reduced search space compared to the original data. Using invert=True also speeds things up a little compared to explicitly negating the resulting mask.

A version of np.isin that relies on the fact that the data is already sorted could be made using np.searchsorted like this:

def isin_sorted(a, i, invert=False):
    ind = np.searchsorted(a, i)
    ind = ind[a[ind.clip(max=a.size)] == i]
    if invert:
        mask = np.ones(a.size, dtype=bool)
        mask[ind] = False
    else:
        mask = np.zeros(a.size, dtype=bool)
        mask[ind] = True
    return mask

You could use this version in place of np.isin, after calling np.unique, which always returns a sorted array. For sufficiently large vec and exclusion lists, it will be more efficient:

result = i[isin_sorted(u, [51, 52], invert=True)]
user_29887300 
import numpy as np
vec = np.array([2, 2, 2, 51, 51, 52, 52, 14, 14, 14, 51, 51, 52, 52])

first_occurrence = []
for x in np.unique(vec):
    if x not in [51,52]:
        first_occurrence.append(np.argmax(x==vec))

argmax finds the index of the first occurrence of the maximum (i.e. True) in the boolean array x==vec. As x is from vec it is guaranteed that there is minimum one True value.

Performance depends on the size of vec and on how many values to find. This simple loop method (blue) outperforms the accepted answer (green and orange) for larger arrays, especially for small numbers of values to find like in the example (for the given toy example it's in fact 1.7 times faster) (source). enter image description here enter image description here It turns out the using unique with index=True is relatively slow, another factor for larger arrays is the memory allocation for the mask.

user_39443220
Related