Numpy - 索引多维数组的一维答案

【问题标题】：Numpy - Indexing one dimension of a multidimensional arrayNumpy - 索引多维数组的一维
【发布时间】：2020-03-25 17:12:24
【问题描述】：

我有一个像这样的 numpy 数组，形状为 (6, 2, 4)：

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)，但没有运气。

【问题讨论】：

标签： arrays python-3.x numpy indexing

【解决方案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]])

或使用显式 integer-array 索引 -

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]])

【讨论】：