从 3D numpy 数组的每个 2D 切片中获取不同的列答案

【问题标题】：Taking different columns from each 2D slice of a 3D numpy array从 3D numpy 数组的每个 2D 切片中获取不同的列
【发布时间】：2020-05-07 18:51:08
【问题描述】：

假设以下 3D numpy 数组：

array([[[4, 1, 3, 5, 0, 1, 5, 4, 3],
        [2, 3, 3, 2, 1, 0, 5, 5, 4],
        [5, 3, 0, 2, 2, 2, 5, 3, 2],
        [0, 3, 1, 0, 2, 4, 1, 1, 5],
        [2, 0, 0, 1, 4, 0, 3, 5, 3]],

       [[2, 2, 4, 1, 3, 4, 1, 1, 5],
        [2, 2, 3, 5, 5, 4, 0, 2, 0],
        [4, 0, 5, 3, 1, 3, 1, 1, 1],
        [4, 5, 0, 0, 5, 3, 3, 2, 4],
        [0, 3, 4, 5, 4, 5, 4, 2, 3]],

       [[1, 3, 2, 2, 0, 4, 5, 0, 2],
        [5, 0, 5, 2, 3, 5, 5, 3, 1],
        [0, 5, 3, 2, 2, 0, 4, 2, 3],
        [4, 4, 0, 3, 2, 1, 5, 3, 0],
        [0, 0, 2, 4, 0, 5, 2, 0, 0]]])

给定一个列表[3, 4, 8]，

是否可以在不使用 for 循环的情况下对给定的张量进行切片？

例如取[0, :, :]的第3列，[1, :, :]的第4列和[2, :, :]的第8列得到：

array([[5, 2, 2, 0, 1],
       [3, 5, 1, 5, 4],
       [2, 1, 3, 0, 0]])

【问题讨论】：

arr[ [0,1,2],:, [3,4,8] ]

标签： numpy multidimensional-array numpy-ndarray array-broadcasting numpy-slicing

【解决方案1】：

这是np.take_along_axis 的一种方式-

In [73]: idx = np.array([3,4,8])

# a is input array
In [72]: np.take_along_axis(a,idx[:,None,None],axis=2)[:,:,0]
Out[72]: 
array([[5, 2, 2, 0, 1],
       [3, 5, 1, 5, 4],
       [2, 1, 3, 0, 0]])

另一个具有显式整数索引的 -

In [79]: a[np.arange(len(idx)),:,idx]
Out[79]: 
array([[5, 2, 2, 0, 1],
       [3, 5, 1, 5, 4],
       [2, 1, 3, 0, 0]])

【讨论】：