python - 在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)。