CODEWITHSUNDEEP
pythonarraysnumpyindexing
user1684046
user1684046

Reputation: 1939

Fast compose slices of 3D numpy array as 3D subarray

I have a 3D numpy array x. I want to take a subset of each slice on axis 0 (each subset is the same shape, but with start and end indices that may be different for each slice) and compose these into a separate 3D numpy array. I can achieve this with

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

but 1) is there any way to do this in a single operation using fancy indexing, and 2) will doing so speed things up?

Upvotes: 2

Views: 257

Answers (1)

Divakar
Divakar

Reputation: 221514

We can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding 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)

Alternatively, a compact version leveraging broadcasting that generates all the required indices and then indexing into the input array, would be -

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

Upvotes: 1

Related Questions