我需要模拟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]))
是否有一个函数可以给我线性索引而不自己计算它们?
答案 0 :(得分:7)
>>> 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]
答案 1 :(得分:3)
最简单的解决方案是在调用nonzero()
之前展平数组:
>>> a.ravel().nonzero()
(array([ 0, 15]),)
答案 2 :(得分:1)
如果您安装了matplotlib
它可能已经存在find
模块中的matplotlib.mlab
,以及其他一些与matlab兼容的功能。是的,它的实现方式与flatnonzero
相同。