numpy，从3d中提取特定行

Question

是否可以从3d数组中提取某些行。

例如，

a = np.arange(27).reshape(3,3,3)

然后，

a = array([[[0,1,2],

            [3,4,5],

            [6,7,8]],

           [[9,10,11],

            [12,13,14],

            [15,16,17]],

            [18,19,20],

            [21,22,23],

            [24,25,26]]])

如果我有随机索引

array([0,1,1])

如何提取与上述索引相对应的某些行，

这意味着结果是

结果=

array([[0,1,2],

       [12,13,14],

       [21,22,23]])

它的维度现在是2d。

Answer 1

a = np.arange(27).reshape(3,3,3)
b = np.array([0,1,1])
print(a[np.arange(len(a)),b])

好