7

I have a 3D numpy array data and another array pos of indexes (an index is a numpy array on its own, which makes the latter array a 2D array):

import numpy as np
data = np.arange(8).reshape(2, 2, -1)
#array([[[0, 1],
#    [2, 3]],
#
#  [[4, 5],
#    [6, 7]]])

pos = np.array([[1, 1, 0], [0, 1, 0], [1, 0, 0]])
#array([[1, 1, 0],
#       [0, 1, 0],
#       [1, 0, 0]])

I want to select and/or mutate the elements from data using the indexes from pos. I can do the selection using a for loop or a list comprehension:

[data[tuple(i)] for i in pos]
#[6, 2, 4]
data[[i for i in pos.T]]
#array([6, 2, 4])

But this does not seem to be a numpy way. Is there a vectorized numpy solution to this problem?

Improve this question
5
  • Do you care if the method returns a view vs a copy? Commented May 16, 2018 at 4:55
  • Not technically an exact dupe. But all relevant info is there. Let me know if you want it reopened. Commented May 16, 2018 at 5:05
  • If you're not satisfied with the answer or the dupe, ping either one of us here. We'll be happy to reopen and slap a bounty on this to incentivise someone smarter to come up with a solution that doesn't exist yet ;-) Commented May 16, 2018 at 5:17
  • 1
    @piRSquared Found an exact dupe. Commented May 16, 2018 at 7:18
  • 1
    @Divakar thank you Commented May 16, 2018 at 7:19

1 Answer 1

Reset to default
3

You can split pos into 3 separate arrays and index, like so—

>>> i, j, k = pos.T
>>> data[i, j, k]
array([6, 2, 4])

Here, the number of columns in pos correspond to the depth of data. As long as you're dealing with 3D matrices, getting i, j, and k well never get more complicated than this.

On python-3.6+, you can shorten this to—

>>> data[[*pos.T]]
array([6, 2, 4])
Improve this answer
Sign up to request clarification or add additional context in comments.

9 Comments

Ouch, downvoter please let me know if there's anything semantically wrong with this answer (from what I can tell, this gives OP what they're expecting). I'd like to correct it! Thanks.
In a sense, you are doing [i for i in pos.T], except that your solution works only for 3 items. My concern is that the list comprehension seems unnecessary here. Isn't there a native numpy way to reorganize the array?
tensorflow has tf.gather_nd. Is there not an analog to numpy?
Yeah. I downvoted this answer and started to post my own using take and then realized it didn't work and removed my downvote. Sorry Coldspeed
@MFisherKDX Well, if OP doesn't mind a python-3.6 solution, I've added it.
|

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.