Search code examples
pythonmatrix-indexing

Obtaining row indices of a matrix that satisfies a condition in Python

I have been trying to obtain all the row indices of the matrix (A) that contains the element nd.

The size of A is 4M*4, it takes ~ 12 sec for this operation.

Link to the file : 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]]

Is there a faster way to accomplish this ?

Solution

  • Since you are using numpy anyway, this speeds up a lot with some more numpy functions. Your current method on my system:

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

    Instead, this is about 40x faster:

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

    You can also look at np.in1d which is similar to the A == nd I have used above, but can compare with lists (something like A == nd1 or nd2 or nd3).