我有一个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解决方案来解决此问题?
您可以分为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])
本文收集自互联网,转载请注明来源。
如有侵权,请联系 [email protected] 删除。
我来说两句