从索引数组访问 numpy 元素答案

【问题标题】：Access numpy elements from an array of indexes从索引数组访问 numpy 元素
【发布时间】：2018-10-05 11:25:06
【问题描述】：

我有一个维度为 n 的数组 a 和一个维度为 n-1 的数组 b。 b 的最后一个轴的值对应于我要从数组 a 中提取的值的索引，位于维度为 n-1 的数组 res 中。

例如 n=2 ：

a = np.array([[1,  2,  3,  4],
              [5,  6,  7,  8],
              [9, 10, 11, 12])
b = np.array([1,3,0])

我愿意

res = [a[1], a[3], a[0]]

# i.e. res = [2, 8, 9]

是否有一种功能可以有效地完成它，具有更多的维度？我知道我可以使用for 循环，但我希望有更高效的方法。

编辑：

在 n=3 的情况下，让 a 的形状为 (2,2,3)。

那么b 和res 的形状为(2,2)：

a = np.array([[[ 1, 2, 3],
               [ 4, 5, 6]],
              [[ 7, 8, 9],
               [10,11,12]]])

b = np.array([[0,2],
              [1,2]]

# res = np.array([[1,6],
#                 [8,12]])

【问题讨论】：

相关：2d案例回答here.
新的np.take_along_axis 会申请吗？
Tkanks 完美运行！我在正确的时间需要它；）

标签： python arrays numpy indexing

【解决方案1】：

最新的numpy（1.15）新增了take_along_axis功能：

In [36]: np.take_along_axis(a, b[:,None], 1)
Out[36]: 
array([[2],
       [8],
       [9]])

它使用辅助函数来构造索引元组：

In [37]: np.lib.shape_base._make_along_axis_idx((3,4), b[:,None], 1)
Out[37]: 
(array([[0],
        [1],
        [2]]), 
 array([[1],
        [3],
        [0]]))

在此之前，我（和其他人）会建议：

In [38]: a[np.arange(3), b]
Out[38]: array([2, 8, 9])

本质上是一样的（除了增加的维度）。正如take_along_axis 文档显示的那样，这旨在将argsort 的结果沿轴获取。

对于更高维度的情况：

In [39]: a1 = np.array([[[ 1, 2, 3],
    ...:                [ 4, 5, 6]],
    ...:               [[ 7, 8, 9],
    ...:                [10,11,12]]])
    ...: b1 = np.array([[0,2],
    ...:               [1,2]])               
In [40]: a1.shape
Out[40]: (2, 2, 3)
In [41]: b1.shape
Out[41]: (2, 2)

In [42]: np.take_along_axis(a1, b1[...,None], -1)
Out[42]: 
array([[[ 1],
        [ 6]],

       [[ 8],
        [12]]])

In [45]: np.lib.shape_base._make_along_axis_idx(a1.shape, b1[...,None], 2)
Out[45]: 
(array([[[0]],

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

        [[1],
         [2]]]))
In [46]: [i.shape for i in _]
Out[46]: [(2, 1, 1), (1, 2, 1), (2, 2, 1)]

同样，等效的自己动手索引：

In [48]: a1[np.arange(2)[:,None], np.arange(2)[None,:], b1]
Out[48]: 
array([[ 1,  6],
       [ 8, 12]])

一旦您了解了数组广播以及它如何应用于索引，这里的概念就不难了。但是take_along_axis 可能会使应用它们更容易。它在某种意义上是np.ix_的延伸。

In [50]: np.ix_(np.arange(2), np.arange(2), np.arange(3))
Out[50]: 
(array([[[0]],

        [[1]]]), array([[[0],
         [1]]]), array([[[0, 1, 2]]]))
In [51]: [i.shape for i in _]
Out[51]: [(2, 1, 1), (1, 2, 1), (1, 1, 3)]
In [55]: a1[(*np.ix_(np.arange(2), np.arange(2)),b1)]
Out[55]: 
array([[ 1,  6],
       [ 8, 12]])

【讨论】：