Numpy：通过z-index数组索引3d数组

Question

说我有一个数组

a = np.zeros((3, 3, 3))

和z-index数组

z = np.random.randint(0, 3, (3, 3))

说z是

array([[1, 0, 2],
       [2, 2, 1],
       [1, 1, 0]])

现在我想用坐标选择a的值（从z的左上角开始并逐行遍历数组。列式也可以。）（0,0 ， 1 ），（0,1， 0 ），（0,2， 2 ），....粗体值来自z数组。

Answer 1

我认为这可以满足您的需求。

import numpy as np 

a = np.arange(3*3*3).reshape(3,3,3)
z = np.array([[1, 0, 2],
       [2, 2, 1],
       [1, 1, 0]])
i = np.arange(a.shape[0]).repeat(a.shape[0]).reshape(a.shape[0], a.shape[1])
j = i.T
#Should do the same as this. possible more efficient, did not test.
#i, j = np.indices(a[...,0].shape)
print a[i,j,z]