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

【问题标题】：Fast compose slices of 3D numpy array as 3D subarray将 3D numpy 数组的切片快速组合为 3D 子数组
【发布时间】：2018-09-30 17:15:53
【问题描述】：

我有一个 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) 这样做会加快速度吗？

【问题讨论】：

开始和结束之间的差异在所有元素中是否总是相同的，比如这里的 2？
是的，没错

标签： python arrays numpy indexing

【解决方案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)]

【讨论】：