在 Python 中使用高级索引进行 nd 数组收缩的简单方法

Question

我知道一定有一种优雅的方法可以使用高级索引来做到这一点，我就是想不出来。

假设我有 (2,3,4) 数组

x = array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

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

和 (4,) 数组 y = array([1,0,1,1])

获得(3,4)数组最优雅的方法是什么

z = np.zeros((3,4))
for ii in xrange(3):
    for jj in xrange(4):
         z[ii,jj] = x[y[jj],ii,jj]

生产？

Answer 1

In [490]: x[y,:,np.arange(4)]                                                                  
Out[490]: 
array([[12, 16, 20],
       [ 1,  5,  9],
       [14, 18, 22],
       [15, 19, 23]])

我们需要转置这个。通过混合使用基本索引和高级索引，切片维度已放在最后：

In [491]: x[y,:,np.arange(4)].T                                                                
Out[491]: 
array([[12,  1, 14, 15],
       [16,  5, 18, 19],
       [20,  9, 22, 23]])

（在一些 SO 中记录并讨论了 basic/advanced 怪癖。）

或使用高级索引：

In [492]: x[y,np.arange(3)[:,None],np.arange(4)]                                               
Out[492]: 
array([[12,  1, 14, 15],
       [16,  5, 18, 19],
       [20,  9, 22, 23]])

在 Python 中使用高级索引进行 nd 数组收缩的简单方法

Easy way to do nd-array contraction using advanced indexing in Python

arrays

numpy

advanced-indexing