获取满足Python条件的矩阵的行索引

Question

我一直在尝试获取包含元素nd的矩阵（A）的所有行索引。

A的大小为4M * 4，此操作需要约12秒。

链接到文件：data

# Download the file named elementconnectivity before running this small script 


A=np.loadtxt('elementconnectivity') 
A=A.astype(int)
nd=1103
v=[i for i in range(len(A)) if nd in A[i]]

有没有更快的方法来实现这一目标？

Answer 1

因为你无论如何都在使用numpy，所以使用更多numpy函数会加速很多。您系统上的当前方法：

%timeit v=[i for i in range(len(A)) if nd in A[i]]
1 loop, best of 3: 4.23 s per loop

相反，这大约快了40倍：

%timeit v = np.where(np.sum((A == nd), axis=1) > 0)
10 loops, best of 3: 92.2 ms per loop

您还可以查看与我上面使用的np.in1d类似的A == nd，但可以与列表进行比较（例如A == nd1或nd2或nd3）。

Answer 2

我认为更好的方法是使用迭代器而不是列表

idxs = (i for i in xrange(len(A)) if nd in A[i])
idxs.next()