arrays - Numpy - 索引多维数组的一维

Question

我有一个形状为 (6, 2, 4) 的 numpy 数组:

x = np.array([[[0, 3, 2, 0],
               [1, 3, 1, 1]],

              [[3, 2, 3, 3],
               [0, 3, 2, 0]],

              [[1, 0, 3, 1],
               [3, 2, 3, 3]],

              [[0, 3, 2, 0],
               [1, 3, 2, 2]],

              [[3, 0, 3, 1],
               [1, 0, 1, 1]],

              [[1, 3, 1, 1],
               [3, 1, 3, 3]]])

我有这样的 choices 数组:

choices = np.array([[1, 1, 1, 1],
                    [0, 1, 1, 0],
                    [1, 1, 1, 1],
                    [1, 0, 0, 0],
                    [1, 0, 1, 1],
                    [0, 0, 0, 1]])

如何使用 choices 数组仅索引大小为 2 的中间维度，并以最有效的方式获取形状为 (6, 4) 的新 numpy 数组可能吗？

结果是这样的:

[[1 3 1 1]
 [3 3 2 3]
 [3 2 3 3]
 [1 3 2 0]
 [1 0 1 1]
 [1 3 1 3]]

我已经尝试通过 x[:, choices, :] 来做到这一点，但这并没有返回我想要的。我也尝试过 x.take(choices, axis=1) 但没有成功。

Answer 1

使用np.take_along_axis沿第二个轴进行索引-

In [16]: np.take_along_axis(x,choices[:,None],axis=1)[:,0]
Out[16]: 
array([[1, 3, 1, 1],
       [3, 3, 2, 3],
       [3, 2, 3, 3],
       [1, 3, 2, 0],
       [1, 0, 1, 1],
       [1, 3, 1, 3]])

或者使用显式整数数组索引-

In [22]: m,n = choices.shape

In [23]: x[np.arange(m)[:,None],choices,np.arange(n)]
Out[23]: 
array([[1, 3, 1, 1],
       [3, 3, 2, 3],
       [3, 2, 3, 3],
       [1, 3, 2, 0],
       [1, 0, 1, 1],
       [1, 3, 1, 3]])