python - 沿 2 个轴的 numpy `take`

Question

我有一个 3D 数组 a 数据和一个 2D 数组 b 索引。我需要使用 b 中的索引沿第 3 轴获取 a 的子数组。我可以像这样用 take 来做到这一点:

a = np.arange(24).reshape((2,3,4))
b = np.array([0,2,1,3]).reshape((2,2))
np.array([np.take(a_,b_,axis=1) for (a_,b_) in zip(a,b)])

我可以在没有列表理解的情况下使用一些花哨的索引来做到这一点吗？我担心效率，所以如果花式索引在这种情况下效率不高，我想知道。

编辑 我试过的第一件事是 a[[0,1],:,b] 但它没有给出我需要的子数组

Answer 1

In [317]: a
Out[317]: 
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]]])
In [318]: a = np.arange(24).reshape((2,3,4))
     ...: b = np.array([0,2,1,3]).reshape((2,2))
     ...: np.array([np.take(a_,b_,axis=1) for (a_,b_) in zip(a,b)])
     ...: 
Out[318]: 
array([[[ 0,  2],
        [ 4,  6],
        [ 8, 10]],

       [[13, 15],
        [17, 19],
        [21, 23]]])

因此，您需要第一个 block 中的 0 和 2 列，以及第二个 block 中的 1 和 3 列。

制作一个形状与b相匹配的c，并体现这一观察结果

In [319]: c=np.array([[0,0],[1,1]])
In [320]: c
Out[320]: 
array([[0, 0],
       [1, 1]])
In [321]: b
Out[321]: 
array([[0, 2],
       [1, 3]])

In [322]: a[c,:,b]
Out[322]: 
array([[[ 0,  4,  8],
        [ 2,  6, 10]],

       [[13, 17, 21],
        [15, 19, 23]]])

这是正确的数字，但不是正确的形状。

可以使用列向量代替 c。

In [323]: a[np.arange(2)[:,None],:,b]  # or a[[[0],[1]],:,b]
Out[323]: 
array([[[ 0,  4,  8],
        [ 2,  6, 10]],

       [[13, 17, 21],
        [15, 19, 23]]])

至于形状，我们可以转置最后两个轴

In [324]: a[np.arange(2)[:,None],:,b].transpose(0,2,1)
Out[324]: 
array([[[ 0,  2],
        [ 4,  6],
        [ 8, 10]],

       [[13, 15],
        [17, 19],
        [21, 23]]])

这个转置是必需的，因为我们在两个索引数组之间有一个切片，混合了基本索引和高级索引。它被记录在案，但从未如此令人费解。它将切片维度 (3) 放在最后，我们必须将它转置回去。

漂亮的小索引拼图!

本次进阶/基础转置的最新问题及解释:

Indexing numpy multidimensional arrays depends on a slicing method