在 numpy 中获取线性化索引答案

【问题标题】：Getting linearized indices in numpy在 numpy 中获取线性化索引
【发布时间】：2023-04-05 05:59:01
【问题描述】：

我需要模拟 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]))

有没有一个函数可以给我线性索引而不自己计算它们？

【问题讨论】：

标签： matlab numpy

【解决方案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]

【讨论】：

【解决方案2】：

最简单的解决方案是在调用nonzero()之前将数组展平：

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

【讨论】：

【解决方案3】：

如果您安装了matplotlib，它可能已经在matplotlib.mlab 模块中存在（即find），以及一些旨在与matlab 兼容的其他功能。是的，它的实现方式与flatnonzero 相同。

【讨论】：