用索引的numpy数组索引numpy数组

Question

我有一个3D numpy数组data和另一个索引数组pos（索引本身就是一个numpy数组，这使得后一个数组成为2D数组）：

import numpy as np
data = np.arange(8).reshape(2, 2, -1)
#array([[[0, 1],
#    [2, 3]],
#
#  [[4, 5],
#    [6, 7]]])

pos = np.array([[1, 1, 0], [0, 1, 0], [1, 0, 0]])
#array([[1, 1, 0],
#       [0, 1, 0],
#       [1, 0, 0]])

我想使用pos中的索引选择和/或改变data中的元素。我可以使用for循环或列表理解进行选择：

^{2}$

但这似乎不是一个裸体的方式。有没有一个向量化的numy解决方案来解决这个问题？在

Answer 1

您可以将pos拆分为3个单独的数组和索引，如下-

>>> i, j, k = pos.T
>>> data[i, j, k]
array([6, 2, 4])

这里，pos中的列数对应于data的深度。只要你在处理3D矩阵，得到i，j，和{}的复杂程度永远不会超过这个。在

在python-3.6+上，可以将其缩短为-

^{pr2}$