Getting linearized indexes in numpy

Question

I need to emulate the MATLAB find function, which returns linear indices for nonzero array elements. For instance:

 >> a = zeros(4,4) a = 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 >> a(1,1) = 1 >> a(4,4) = 1 >> find(a) ans = 1 16 

numpy has a similar nonzero function, but returns a tuple of index arrays. For instance:

 In [1]: from numpy import * In [2]: a = zeros((4,4)) In [3]: a[0,0] = 1 In [4]: a[3,3] = 1 In [5]: a Out[5]: array([[ 1., 0., 0., 0.], [ 0., 0., 0., 0.], [ 0., 0., 0., 0.], [ 0., 0., 0., 1.]]) In [6]: nonzero(a) Out[6]: (array([0, 3]), array([0, 3])) 

Is there a function that gives me linear indices without calculating them myself?

Answer 1

numpy you covered:

 >>> np.flatnonzero(a) array([ 0, 15]) 

Inside, he does exactly what Sven Marnach suggested.

 >>> print inspect.getsource(np.flatnonzero) def flatnonzero(a): """ Return indices that are non-zero in the flattened version of a. This is equivalent to a.ravel().nonzero()[0]. [more documentation] """ return a.ravel().nonzero()[0] 

Answer 2

The simplest solution is to smooth the array before calling nonzero() :

 >>> a.ravel().nonzero() (array([ 0, 15]),) 

Answer 3

If you have matplotlib installed, it may already exist ( find ) in the matplotlib.mlab module, as well as some other functions designed for compatibility with matlab. And yes, it is implemented in the same way as flatnonzero .