CODEWITHSUNDEEP
pythonmatrix-indexing
Mechanician
Mechanician

Reputation: 545

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 ?

Upvotes: 1

Views: 191

Answers (2)

VBB
VBB

Reputation: 1325

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).

Upvotes: 1

Yunier Rojas
Yunier Rojas

Reputation: 79

I thinks a better way to accomplish this is to use iterators instead of lists

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

Upvotes: 0

Related Questions