Getting linearized indices in numpy

matlab, numpy

Solution

numpy has you covered:

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

Internally it's doing 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]

Problem

I need to emulate the MATLAB function `find`, which returns the linear indices for the nonzero elements of an array. For example: ``` >> 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 the similar function `nonzero`, but it returns a tuple of index arrays. For example: ``` 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 the linear indices without calculating them myself?

Original source