python - 如何沿给定轴取元素，由它们的索引给出？

Question

我有一个 3D 数组，我需要在最后一个轴上“挤压”它，这样我就得到了一个 2D 数组。我需要按照以下方式进行。对于前两个维度的每个索引值，我知道第三个维度的索引值应该从哪里获取。

例如，我知道如果 i1 == 2 和 i2 == 7 那么 i3 == 11。这意味着 out[2,7] = inp[2,7,11]。从前两个维度到第三个维度的映射在另一个二维数组中给出。换句话说，我有一个数组，其中 2,7 的值是 11。

所以，我的问题是如何组合这两个数组(3D 和 2D)以获得输出数组(2D)。

Answer 1

In [635]: arr = np.arange(24).reshape(2,3,4)
In [636]: idx = np.array([[1,2,3],[0,1,2]])


In [637]: I,J = np.ogrid[:2,:3]
In [638]: arr[I,J,idx]
Out[638]: 
array([[ 1,  6, 11],
       [12, 17, 22]])
In [639]: arr
Out[639]: 
array([[[ 0,  1,  2,  3],   # 1
        [ 4,  5,  6,  7],   # 6
        [ 8,  9, 10, 11]],  # ll

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])

I,J一起广播选择一组(2,3)值，匹配idx:

In [640]: I
Out[640]: 
array([[0],
       [1]])
In [641]: J
Out[641]: array([[0, 1, 2]])

---

这是更简单的 2d 问题的 3d 泛化 - 从每一行中选择一个项目:

In [649]: idx
Out[649]: 
array([[1, 2, 3],
       [0, 1, 2]])
In [650]: idx[np.arange(2), [0,1]]
Out[650]: array([1, 1])

事实上，我们可以将 3d 问题转换为 2d 问题:

In [655]: arr.reshape(6,4)[np.arange(6), idx.ravel()]
Out[655]: array([ 1,  6, 11, 12, 17, 22])

---

概括原始案例:

In [55]: arr = np.arange(24).reshape(2,3,4)                                     
In [56]: idx = np.array([[1,2,3],[0,1,2]])                                      
In [57]: IJ = np.ogrid[[slice(i) for i in idx.shape]]                           
In [58]: IJ                                                                     
Out[58]: 
[array([[0],
        [1]]), array([[0, 1, 2]])]
In [59]: (*IJ,idx)                                                              
Out[59]: 
(array([[0],
        [1]]), array([[0, 1, 2]]), array([[1, 2, 3],
        [0, 1, 2]]))
In [60]: arr[_]                                                                 
Out[60]: 
array([[ 1,  6, 11],
       [12, 17, 22]])

关键在于将数组的 IJ 列表与 idx 组合起来以创建新的索引元组。如果 idx 不是最后一个索引，构建元组会有点困惑，但它仍然是可能的。例如

In [61]: (*IJ[:-1],idx,IJ[-1])                                                  
Out[61]: 
(array([[0],
        [1]]), array([[1, 2, 3],
        [0, 1, 2]]), array([[0, 1, 2]]))
In [62]: arr.transpose(0,2,1)[_]                                                
Out[62]: 
array([[ 1,  6, 11],
       [12, 17, 22]])

如果将 arr 转置为 idx 维度更容易，则放在最后。关键是索引操作采用索引数组的元组，这些数组相互广播以选择特定项目。这就是 ogrid 所做的，创建与 idx 一起使用的数组。