I am trying to get indices of elements that satisfies a condition in numpy array. For instance:
import numpy as np
ar = np.array([1, 1, 2, 1, 3, 4, 1])
np.where(ar==1)[0]
will return:
[0, 1, 3, 6]
Now if I want the condition to be a list of values like:
np.where(ar in [1, 2])[0]
I would like the result to be:
[0, 1, 2, 3, 6]
How can I achieve this using numpy? Of course I can iterate over my list of conditions and concatenate the results
indices = np.concatenate([np.where(ar==cond)[0] for cond in [1, 2])
But I would like to know if there's a more direct numpy-like way of doing this.
