MATLAB-style find() function in Python MATLAB-style find() function in Python python python

MATLAB-style find() function in Python


in numpy you have where :

>> import numpy as np>> x = np.random.randint(0, 20, 10)>> xarray([14, 13,  1, 15,  8,  0, 17, 11, 19, 13])>> np.where(x > 10)(array([0, 1, 3, 6, 7, 8, 9], dtype=int64),)


You can make a function that takes a callable parameter which will be used in the condition part of your list comprehension. Then you can use a lambda or other function object to pass your arbitrary condition:

def indices(a, func):    return [i for (i, val) in enumerate(a) if func(val)]a = [1, 2, 3, 1, 2, 3, 1, 2, 3]inds = indices(a, lambda x: x > 2)>>> inds[2, 5, 8]

It's a little closer to your Matlab example, without having to load up all of numpy.


Or use numpy's nonzero function:

import numpy as npa    = np.array([1,2,3,4,5])inds = np.nonzero(a>2)a[inds] array([3, 4, 5])