CODEWITHSUNDEEP
pythonnumpy
rahul003
rahul003

Reputation: 165

Is there a faster method to do this sort of indexing in numpy?

I have numpy arrays a and b as below.

a.shape = (100,10000)
b.shape = (100,)

Need a c of shape (100,) such that,

for i in range(0,len(b)):
   c[i] = a[i,b[i]]

I want a faster way of doing the above. For loops can't be used in the scenario I'm using.. It works, but still expensive. Is there a more efficient way of doing it?

Upvotes: 3

Views: 64

Answers (1)

DrV
DrV

Reputation: 23490

Yes, you could do:

import numpy as np

c = a[np.arange(len(b)), b]

That is about as fast as it gets. In the following code

import numpy as np

a = np.random.random((1000000, 10))
b = np.random.randint(0,10,1000000)
c = a[np.arange(len(b)), b]

the last operation takes 31 ms (by IPython's %timeit), i.e. 31 ns per indexing.

Upvotes: 3

Related Questions