Question

我有一个尺寸为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]])

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

从索引数组访问numpy元素

1 个答案: