Numpy：通过另一个数组组合数组列表（np.choose 替代）

Question

我有一个 numpy 数组列表，每个数组的形状都相同。比方说：

a = [np.array([[1, 2, 3],
               [4, 5, 6],
               [7, 8, 9]]),
     np.array([[11, 12, 13],
               [14, 15, 16],
               [17, 18, 19]]),
     np.array([[99, 98, 97],
               [96, 95, 94],
               [93, 92, 91]])]

我还有另一个相同形状的数组，它给出了我想从中获取元素的列表索引：

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

我想要得到的是以下内容：

np.array([[1, 2, 13],
          [96, 15, 6],
          [93, 18, 91]])

有一个简单的解决方案效果很好：

np.choose(b, a)

但这仅限于最多 32 个数组。但就我而言，我必须组合更多的数组（超过 100 个）。所以我需要另一种方法。

我想，它必须是关于高级索引或者np.take 方法的东西。所以可能，第一步是a = np.array(a)，然后是a[np.arange(a.shape[0]), b]。但我没有得到它的工作。

有人可以帮忙吗？ :)

Answer 1

您可以尝试使用np.ogrid。基于this 的回答。当然，您必须先将a 转换为 NumPy 数组

i, j = np.ogrid[0:3, 0:3]
print (a[b, i, j])

# array([[ 1,  2, 13],
#        [96, 15,  6],
#        [93, 18, 91]])

Answer 2

In [129]: a = [np.array([[1, 2, 3], 
     ...:                [4, 5, 6], 
     ...:                [7, 8, 9]]), 
     ...:      np.array([[11, 12, 13], 
     ...:                [14, 15, 16], 
     ...:                [17, 18, 19]]), 
     ...:      np.array([[99, 98, 97], 
     ...:                [96, 95, 94], 
     ...:                [93, 92, 91]])]                                        
In [130]: b = np.array([[0, 0, 1], 
     ...:               [2, 1, 0], 
     ...:               [2, 1, 2]])                                             
In [131]:                                                                       
In [131]: A = np.array(a)                                                       
In [132]: A.shape                                                               
Out[132]: (3, 3, 3)

您想使用b 来索引第一个维度。对于其他维度，您需要使用b 广播的索引，即列向量和行向量：

In [133]: A[b, np.arange(3)[:,None], np.arange(3)]                              
Out[133]: 
array([[ 1,  2, 13],
       [96, 15,  6],
       [93, 18, 91]])

有各种方便的函数来创建这些数组，例如

In [134]: np.ix_(range(3),range(3))                                             
Out[134]: 
(array([[0],
        [1],
        [2]]), array([[0, 1, 2]]))

和ogrid 如另一个答案中所述。

这是一个相对较新的功能，也可以完成这项工作：

In [138]: np.take_along_axis(A, b[None,:,:], axis=0)                              
Out[138]: 
array([[[ 1,  2, 13],
        [96, 15,  6],
        [93, 18, 91]]])

在将b 调整为正确之前，我不得不考虑一下。