Python 使用另一个数组作为深度对 numpy 数组进行切片

Question

我有一个维度为 (x = 2, y = 2, z = 3) 的 3 维 numpy 数组，如下所示

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

我想从 a 的每个 (x, y) 元素中获取前 N 个元素，其中 N 在另一个大小为 x,y 的数组中定义。例如

b = [[1,2],
     [0,0]]

结果应该是

c = [[[0,1],[3,4,5]],
    [[6],[9]]

如何在没有循环的情况下执行此操作？

Answer 1

您可以在 a 中给定范围内的元素创建一个掩码，并在需要对 a 进行操作时使用此掩码。如果您希望将掩码和数组封装在一个实体中，请查看 numpy.ma 模块。

a = np.array([[[0,1,2],[3,4,5]],
             [[6,7,8],[9,10,11]]])

b = np.array([[1,2],
              [0,0]])

mask = np.arange(3)[None,None,:] <= b[:,:,None]

a[mask]

输出：

array([0, 1, 3, 4, 5, 6, 9])

但是，如果您希望输出为非同类数组，我想不出比使用循环更好的方法了。

Python 使用另一个数组作为深度对 numpy 数组进行切片

Python slicing numpy array using another array as depth

python

arrays

numpy

slice