1

I have a bidimensional array A and a list of indexes idx, for example : 

A = np.array([[ 1.,  0.,  0.],
             [ 0.,  1.,  0.],
             [ 0.,  0.,  1.],
             [ 0., -1.,  0.],
             [ 0.,  0.,  5.]])

idx = np.array([2, 1, 0, 1, 2])

and I'm trying to select the elements of A indexed by idx along the column axis (in this example : array([0., 1., 0., -1., 5.])). How can I do this without loops ?

Thank you !

Improve this question
3
  • Why do you not want to use loops? - This could easily be solved with a for loop. Commented Aug 3, 2018 at 12:53
  • 2
    @Alfie because numpy is usually extremely faster than for loops. Commented Aug 3, 2018 at 12:57
  • @Alfie As @Luca said, numpy operations are extremely more efficient than loops in Python, and on top of that the matrix A has huge dimensions in my problem. Commented Aug 3, 2018 at 13:09

2 Answers 2

Reset to default
1 
A[np.arange(np.size(idx)), idx]

gives array([ 0., 1., 0., -1., 5.])

From the Advanced Indexing part of the documentation:

When the index consists of as many integer arrays as the array being indexed has dimensions, the indexing is straight forward, but different from slicing. [...] This is best understood with an example.

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

Comments

0

Indexing 2D numpy arrays can be a bit confusing.

You need: A[np.arange(0, A.shape[0]), idx]

Improve this answer

Comments

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.