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

你可以分成pos3个单独的数组和索引，像这样\xe2\x80\x94

\n\n

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

\n\n

这里， 的列数pos对应于 的深度data。只要您处理 3D 矩阵，获取i、j、 以及k就不会比这更复杂了。

\n\n

在 python-3.6+ 上，您可以将其缩短为\xe2\x80\x94

\n\n

>>> data[[*pos.T]]\narray([6, 2, 4])\n

\n