I have a large numpy array of shape (n,m). I need to extract one element from each row, and I have another array of shape (n,) that gives the column index of the element I need. The following code does this, but it requires an explicit loop (in the form of a list comprehension.):

import numpy as np

arr = np.array(range(12))
arr = arr.reshape((4,3))
keys = np.array([1,0,1,2])
#This is the line that I'd like to optimize
answers = np.array([arr[i,keys[i]] for i in range(len(keys))])
# [ 1  3  7 11]

Is there a built-in numpy (or pandas?) function that could do this more efficiently?

  • \$\begingroup\$ arr[np.arange(4), keys] \$\endgroup\$ Nov 1 '15 at 10:54
  • 1
    \$\begingroup\$ It's correct. @GarethRees Please put that in the answer so you don't waste others time looking at something that hasn't been answered yet. \$\endgroup\$ Nov 1 '15 at 14:26

The best way to do this is how @GarethRees suggested in the comments:

>>> arr[np.arange(4), keys]
array([1, 3, 7, 11])

There is another (not as good) solution that is a better alternative to using arr[i, keys[i]] for i in range(len(keys)). Whenever you want both the index and the item from some iterable, you should generally use the enumerate function:

>>> np.array([arr[i, item] for i, item in enumerate(keys)])
array([1, 3, 7, 11])
  • 1
    \$\begingroup\$ Good point about enumerate since it's generally applicable. Thanks to Gareth for providing the initial answer in a comment. I'm accepting this one so that the question doesn't show as unanswered. \$\endgroup\$
    – Max Rosett
    Nov 2 '15 at 5:08

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.