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

Question

假设以下 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]])

Answer 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]])

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

Taking different columns from each 2D slice of a 3D numpy array

numpy

multidimensional-array

numpy-slicing

numpy-ndarray

array-broadcasting