使用一个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循环或列表推导来进行选择：

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

但这似乎不是numpy的方式。是否有一种向量化的numpy解决方案来解决这个问题？

- DYZ

---

您介意该方法返回视图还是副本吗？ - user3483203

不是完全相同的重复。但所有相关信息都在那里。如果您想重新打开它，请告诉我。 - piRSquared

如果你对答案或重复问题不满意，可以在这里@我们中的任何一个。我们很乐意重新开放并添加悬赏以激励更聪明的人提供尚不存在的解决方案;-) - cs95

@piRSquared 找到了一个完全重复的。 - Divakar

@Divakar 谢谢您 - piRSquared

Answer 1

你可以将pos分成三个单独的数组并进行索引，如下所示——

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

在此，pos 中的列数对应于 data 的深度。只要您处理的是三维矩阵，则获得i，j 和 k 永远不���比这更复杂。

在 Python 3.6+ 上，您可以将其缩短为 —

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

- cs95

---

哎呀，罩衣请告诉我这个答案中是否有任何语义错误（据我所知，这可以为 OP 提供他们期望的东西）。 我想纠正它！谢谢。 - cs95

TensorFlow有tf.gather_nd。NumPy没有类似的吗？ - MFisherKDX

是的。我给这个答案投了反对票，然后开始用 take 发布我的答案，但后来意识到它不起作用了，就取消了我的反对票。抱歉，Coldspeed。 - MFisherKDX

@MFisherKDX 如果OP不介意使用Python 3.6的解决方案，我已经添加了它。 - cs95

自 v1.15.0 起，data[[*pos.T]] 已被弃用: "使用除元组以外的任何多维索引都已被弃用。" 请改用 data[tuple(pos.T)]。 - John