Быстрое создание фрагментов трехмерного массива в виде 3D-массива

Question

У меня есть трехмерный массив NumPy x.Я хочу взять подмножество каждого среза по оси 0 (каждое подмножество имеет одну и ту же форму, но с индексами начала и конца, которые могут отличаться для каждого среза) и объединить их в отдельный трехмерный массив с фрагментами.Я могу добиться этого с помощью

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, чтобы получить раздвижные окна. Дополнительная информация по использованию as_strided на основе 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)]