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

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

[data[tuple(i)] for i in pos]
#[6, 2, 4]
data[[i for i in pos.T]]
#array([6, 2, 4])

但这似乎并不是一种愚蠢的方式。这个问题有没有矢量化的numpy解决方案？

Answer 1

您可以将pos分成3个独立的数组和索引，就像这样 -

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

此处，pos中的列数对应data的深度。只要您处理3D矩阵，获得i，j和k就不会比这更复杂。

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

>>> data[[*pos.T]]
array([6, 2, 4])