使用2维数组索引3维numpy数组

Question

我想创建一个基于numpy 3d数组中的值的numpy 2d数组，使用另一个numpy 2d数组来确定在轴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]]

第二个和第三个轴元素的两个示例结果展示了我使用这两个轴元素时得到的结果。但是我希望得到由 arr_2d 指定的 arr_3d 中的元素。因此...

- 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应该长这样...

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

我尝试使用这条语句获取arr_2d的项，但结果是一个4维数组，而我需要一个2维数组。

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

- epe

Answer 1

使用花式索引和广播后的结果形状与索引数组的形状相同。 您需要为 arr_3d 的每个轴传递 2D 数组。

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]])

- Andy L.

Answer 2

虽然在功能上与其他答案相同，但我更喜欢使用numpy.ogrid来生成（稀疏的）索引数组，以索引其余维度：

import numpy as np

# example data
arr_3d = np.arange(2*3*4).reshape(2, 3, 4)
arr_2d = np.array(([3, 2, 0], [2, 3, 2]))

m, n, _ = arr_3d.shape
ind0, ind1 = np.ogrid[:m, :n]
res_2d = arr_3d[ind0, ind1, arr_2d]
print(res_2d)
# [[ 3  6  8]
#  [14 19 22]]

我们还可以使用带有sparse=True的np.meshgrid，但那样会强制我们创建1D索引范围，而我们可以通过np.ogrid助手避免这种情况。

- Andras Deak -- Слава Україні

Answer 3

In [107]: arr_3d = np.arange(2*3*4).reshape(2,3,4)                                                           
In [108]: arr_2d = np.array(([3,2,0], [2,3,2]))                                                              
In [109]: arr_2d.shape                                                                                       
Out[109]: (2, 3)
In [110]: arr_3d[[[0],[1]],[0,1,2],arr_2d]                                                                   
Out[110]: 
array([[ 3,  6,  8],
       [14, 19, 22]])

< p > [[0],[1]]，[0,1,2] 相互广播以索引与 `arr_2d` 相同大小的 (2,3) 块。

ix_ 可用于构建这些 2 个索引：

In [114]: I,J = np.ix_(range(2), range(3))                                                                   
In [115]: I,J                                                                                                
Out[115]: 
(array([[0],
        [1]]), array([[0, 1, 2]]))
In [116]: arr_3d[I, J, arr_2d]                                                                               
Out[116]: 
array([[ 3,  6,  8],
       [14, 19, 22]])

- hpaulj