0

This question already has an answer here: Index n dimensional array with (n-1) d array 1 answer

I think the answer here is more straightforward (I understand the problem right away). Repetitive posts are great, and I associate them with a common indexing function, but I need time to understand.

I want to extract the data from the 3d array by the index i query along the axis, but I don't know how to write it. 

#data:
arr= np.array([[[ 0.000,  0.200],
                [ 0.000,  0.100],
                [-0.932, -0.073]],

               [[ 0.000,  0.000],
                [-0.932, -0.073],
                [-1.626, -0.900]],

               [[-0.132, -0.073],
                [-1.626, -0.900],
                [-1.802, -0.688]],

               [[-1.626, -0.900],
                [-1.802, -0.688],
                [-3.059, -1.190]]])

# This is the index
idx= np.array([[0, 1],
               [1, 2],
               [0, 0],
               [0, 1]])

# Expected output
    array([[ 0.000,  0.100],
           [-0.932, -0.900],
           [-0.132, -0.073],
           [-1.626, -0.688],

# If it is 2d it is like this:
axis=0:
arr[idx,np.arange(arr.shape[1])]
axis=1:
arr[np.arange(arr.shape[0]),idx]

But I don’t know how to write it on the 3d array, I hope to get help. Thank you.

weidong
  • 159
  • 8

1 Answers1

2

You weren't very clear about which axis idx applied to. I had match up values to see the pattern.

Anyways here's the indexing

In [88]: arr[np.arange(4)[:,None], idx, np.arange(2)]
Out[88]: 
array([[ 0.   ,  0.1  ],
       [-0.932, -0.9  ],
       [-0.132, -0.073],
       [-1.626, -0.688]])

The 2 aranges were written so they broadcast to the same (4,2) shape as idx.

hpaulj
  • 221,503
  • 14
  • 230
  • 353