使用 2D 数组沿 1 个轴对 3D 数组进行索引

Question

假设我有一个形状为 [z, y, x] 的 numpy 数组 a 和另一个数组 b形状：[y, x]。

数组 b 包含沿 z 的索引，我想为每个 y 和 x 从 a 中提取这些索引.

到目前为止，我有以下不雅的方法：

 from_a = np.full_like(b, np.nan)

 for i in range(y):
  for j in range(x): 

   des_ind = b[i,j]

   from_a[i,j] = a[des_ind,i,j]

有没有很好的 Pythonic 方式来做到这一点？

Answer 1

您可以直接使用数组b 作为a 的索引：

import numpy as np
z, y, x = 3, 4, 5
a = np.random.randint(10, (z, y, x))
# array([[[2, 1, 8, 3, 0],  | [[9, 3, 6, 8, 7],  | [[6, 6, 2, 1, 4],
#         [4, 0, 1, 2, 9],  |  [5, 9, 0, 7, 1],  |  [3, 9, 4, 4, 8],
#         [0, 0, 7, 2, 6],  |  [7, 9, 8, 4, 9],  |  [9, 7, 0, 3, 0],
#         [3, 5, 2, 9, 3]], |  [6, 8, 9, 5, 9]], |  [4, 3, 0, 1, 1]]])

b = np.random.randint(z, size=(y, x))
# array([[1, 1, 1, 0, 2],
#        [1, 2, 1, 1, 2],
#        [2, 0, 0, 0, 2],
#        [0, 0, 1, 0, 0]])

from_a = a[b, 
           np.arange(y)[:, np.newaxis], 
           np.arange(x)[np.newaxis, :]]
# array([[9, 3, 6, 3, 4],
#        [5, 9, 0, 7, 8],
#        [9, 0, 7, 2, 0],
#        [3, 5, 9, 9, 3]])

有关 numpy 数组，另请参阅 advanced indexing。