我有以下数组:
a = np.arange(12).reshape((2, 2, 3))
和
b = np.zeros((2, 2))
现在我想用b访问a,s.t。对于索引 i,j,如果 b[i, j] = z,则我们取 a 的第 z 个元素。 对于上面的例子来说,答案应该是 [[0, 3], [6, 9]]。 我觉得这与 np.choose 非常相关,但不知何故无法完全管理它。 你能帮我吗?
最佳答案
可以建议两种方法。
使用显式
范围数组 advanced-indexing
-
m,n = b.shape
out = a[np.arange(m)[:,None],np.arange(n),b.astype(int)]
np.take_along_axis(a,b.astype(int)[...,None],axis=2)[...,0]
示例运行 -
In [44]: a
Out[44]:
array([[[ 0, 1, 2],
[ 3, 4, 5]],
[[ 6, 7, 8],
[ 9, 10, 11]]])
In [45]: b
Out[45]:
array([[0., 0.],
[0., 0.]])
In [46]: m,n = b.shape
In [47]: a[np.arange(m)[:,None],np.arange(n),b.astype(int)]
Out[47]:
array([[0, 3],
[6, 9]])
In [48]: np.take_along_axis(a,b.astype(int)[...,None],axis=2)[...,0]
Out[48]:
array([[0, 3],
[6, 9]])
关于python - 使用数组作为多维数组的索引掩码,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/55605833/