用2d数组索引3d numpy数组 [英] Indexing 3d numpy array with 2d array

问题描述

我想基于一个numpy 3d数组中的值创建一个numpy 2d数组，使用另一个numpy 2d数组确定在轴3中使用哪个元素.

I would like to create a numpy 2d-array based on values in a numpy 3d-array, using another numpy 2d-array to determine which element to use in axis 3.

import numpy as np
#--------------------------------------------------------------------
arr_3d = np.arange(2*3*4).reshape(2,3,4)
print('arr_3d shape=', arr_3d.shape, '\n', arr_3d)
arr_2d = np.array(([3,2,0], [2,3,2]))
print('\n', 'arr_2d shape=', arr_2d.shape, '\n', arr_2d)
res_2d = arr_3d[:, :, 2]
print('\n','res_2d example using element 2 of each 3rd axis...\n', res_2d)
res_2d = arr_3d[:, :, 3]
print('\n','res_2d example using element 3 of each 3rd axis...\n', res_2d)

结果...

arr_3d shape= (2, 3, 4) 
 [[[ 0  1  2  3]
  [ 4  5  6  7]
  [ 8  9 10 11]]

 [[12 13 14 15]
  [16 17 18 19]
  [20 21 22 23]]]

 arr_2d shape= (2, 3) 
 [[3 2 0]
 [2 3 2]]

 res_2d example using element 2 of each 3rd axis...
 [[ 2  6 10]
 [14 18 22]]

 res_2d example using element 3 of each 3rd axis...
 [[ 3  7 11]
 [15 19 23]]

第2个示例结果显示了如果我使用轴3的第2个元素，然后使用第3个元素，则会得到什么.但是我想从arr_2d指定的arr_3d中获得该元素.所以...

The 2 example results show what I get if I use the 2nd and then the 3rd element of axis 3. But I would like to get the element from arr_3d, specified by arr_2d. So...

- res_2d[0,0] would use the element 3 of arr_3d axis 3
- res_2d[0,1] would use the element 2 of arr_3d axis 3
- res_2d[0,2] would use the element 0 of arr_3d axis 3
etc

所以res_2d应该看起来像这样...

So res_2d should look like this...

[[3 6 8]
[14 19 22]]

我尝试使用此行获取arr_2d条目，但结果为4维数组，而我需要2维数组.

I tried using this line to get the arr_2d entries, but it results in a 4-dim array and I want a 2-dim array.

res_2d = arr_3d[:, :, arr_2d[:,:]]

推荐答案

花式索引和广播结果的形状是索引数组的形状.您需要为arr_3d

The shape of the result from fancy index and broadcasting is the shape of the indexing array. You need passing 2d array for each axis of arr_3d

ax_0 = np.arange(arr_3d.shape[0])[:,None]
ax_1 = np.arange(arr_3d.shape[1])[None,:]

arr_3d[ax_0, ax_1, arr_2d]

Out[1127]:
array([[ 3,  6,  8],
       [14, 19, 22]])

这篇关于用2d数组索引3d numpy数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！