Получение линеаризованных индексов в numpy

Question

Мне нужно эмулировать функцию MATLAB find, которая возвращает линейные индексы для ненулевых элементов массива. Например:

>> 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 имеет аналогичную функцию nonzero, но он возвращает кортеж из индексных массивов. Например:

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])) 

Есть ли функция, которая дает мне линейные индексы, не вычисляя их самостоятельно?

Answer 1

NumPy покрывали Вас:

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

Внутренне он делает именно то, что предложил Sven Marnach.

>>> 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

Самое простое решение, чтобы сгладить массив перед вызовом nonzero():

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

Answer 3

Если вы matplotlib установили его, вероятно, уже есть (find, что есть) в matplotlib.mlab модуле, а также некоторых других функций, предназначенных для обеспечения совместимости с MATLAB. И да, это реализовано так же, как flatnonzero.