Efficiently extract values from array using a list of index

Question

Given a 2D NumPy array a and a list of indices stored in index, there must be a way of extracting the values of the list very efficiently. Using a for loop as follows take about 5 ms which seems extremely slow for 2000 elements to extract:

import numpy as np
import time

# generate dummy array 
a = np.arange(4000).reshape(1000, 4) 
# generate dummy list of indices
r1 = np.random.randint(1000, size=2000)
r2 = np.random.randint(3, size=2000)
index = np.concatenate([[r1], [r2]]).T

start = time.time()
result = [a[i, j] for [i, j] in index]
print time.time() - start

How can I increase the extraction speed? np.take does not seem appropriate here because it would return a 2D array instead of a 1D array.

akuiper · Accepted Answer · 2018-04-15 21:07:18Z

2

You can use advanced indexing which basically means extract the row and column indices from the index array and then use it to extract values from a, i.e. a[index[:,0], index[:,1]] -

%timeit a[index[:,0], index[:,1]]
# 12.1 µs ± 368 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit [a[i, j] for [i, j] in index]
# 2.22 ms ± 105 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

answered Apr 15, 2018 at 21:07

akuiper

216k33 gold badges362 silver badges379 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

miradulo · Accepted Answer · 2018-04-15 21:26:03Z

2

Another option would be numpy.ravel_multi_index, which lets you avoid the manual indexing.

np.ravel_multi_index(index.T, a.shape)

answered Apr 15, 2018 at 21:26

miradulo

29.8k7 gold badges86 silver badges97 bronze badges

1 Comment

anderstood Over a year ago

Thank you. It seems to be even faster because I can use indexT = index.T (once for all).

Collectives™ on Stack Overflow

Efficiently extract values from array using a list of index

2 Answers 2

Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related