MATLAB-style find() function in Python
In MATLAB it is easy to find the indices of values that meet a particular condition:
>> a = [1,2,3,1,2,3,1,2,3];
>> find(a > 2) % find the indecies where this condition is true
[3, 6, 9] % (MATLAB uses 1-based indexing)
>> a(find(a > 2)) % get the values at those locations
[3, 3, 3]
What would be the best way to do this in Python?
So far, I have come up with the following. To just get the values:
>>> a = [1,2,3,1,2,3,1,2,3]
>>> [val for val in a if val > 2]
[3, 3, 3]
But if I want the index of each of those values it's a bit more complicated:
>>> a = [1,2,3,1,2,3,1,2,3]
>>> inds = [i for (i, val) in enumerate(a) if val > 2]
>>> inds
[2, 5, 8]
>>> [val for (i, val) in enumerate(a) if i in inds]
[3, 3, 3]
Is there a better way to do this in Python, especially for arbitrary conditions (not just 'val > 2')?
I found functions equivalent to MATLAB 'find' in NumPy but I currently do not have access to those libraries.
Answer
in numpy you have where :
>> import numpy as np
>> x = np.random.randint(0, 20, 10)
>> x
array([14, 13, 1, 15, 8, 0, 17, 11, 19, 13])
>> np.where(x > 10)
(array([0, 1, 3, 6, 7, 8, 9], dtype=int64),)