使用2-D数组切片3-D数组

Question

Assume we have two matrices:

x = np.random.randint(10, size=(2, 3, 3))
idx = np.random.randint(3, size=(2, 3))

The question is to access the element of x using idx , in the way as:

dim1 = x[0, range(0,3), idx[0]]  # slicing x[0] using idx[0]
dim2 = x[1, range(0,3), idx[1]]
res = np.vstack((dim1, dim2))

Is there a neat way to do this?

Answer 1

You can just index it the basic way , only that the size of indexer array has to match. That's what those .reshape s are for:

x[np.array([0,1]).reshape(idx.shape[0], -1), 
  np.array([0,1,2]).reshape(-1,idx.shape[1]),
  idx]
Out[29]: 
array([[ 0.10786251,  0.2527514 ,  0.11305823],
       [ 0.67264076,  0.80958292,  0.07703623]])

Answer 2

Here's another way to do it with reshaping -

x.reshape(-1,x.shape[2])[np.arange(idx.size),idx.ravel()].reshape(idx.shape)

Sample run -

In [2]: x
Out[2]: 
array([[[5, 0, 9],
        [3, 0, 7],
        [7, 1, 2]],

       [[5, 3, 5],
        [8, 6, 1],
        [7, 0, 9]]])

In [3]: idx
Out[3]: 
array([[2, 1, 2],
       [1, 2, 0]])

In [4]: x.reshape(-1,x.shape[2])[np.arange(idx.size),idx.ravel()].reshape(idx.shape)
Out[4]: 
array([[9, 0, 2],
       [3, 1, 7]])