python - 将 3D numpy 数组的切片快速组合为 3D 子数组

Question

我有一个 3D numpy 数组 x。我想在轴 0 上获取每个切片的一个子集(每个子集的形状相同，但每个切片的开始和结束索引可能不同)并将它们组合成一个单独的 3D numpy 数组。我可以用

import numpy as np

x = np.arange(24).reshape((3, 4, 2))
starts = [0, 2, 1]
ends = [2, 4, 3]

np.stack([x[i, starts[i]:ends[i]] for i in range(3)])

但 1) 是否有任何方法可以使用花哨的索引在单个操作中完成此操作，以及 2) 这样做会加快速度吗？

Answer 1

我们可以利用 np.lib.stride_tricks.as_strided基于 scikit-image's view_as_windows获得滑动窗口。 More info on use of as_strided based view_as_windows .

from skimage.util.shape import view_as_windows

L = 2 # ends[0]-starts[0]
w = view_as_windows(x,(1,L,1))[...,0,:,0]
out = w[np.arange(len(starts)), starts].swapaxes(1,2)

或者，利用 broadcasting 的紧凑版本生成所有必需的索引，然后索引到输入数组中，将是 -

x[np.arange(len(starts))[:,None],np.asarray(starts)[:,None] + np.arange(L)]