numpy select multiple ranges from a 1d array

Question

Let's say I have a 1D array of values:

T = np.array([1.3, 8.9, 1.4, 3.2, 4.4, 7.0, 2.0, 6.9]

and I have a list of start indices:

I = np.array([5, 2, 4, 1])

For each start index, I would like to grab m consecutive values from T starting at the start index. So, the final output for m=3 should be a 2D array with:

T = np.array([[7.0, 2.0, 6.9],
              [1.4, 3.2, 4.4],
              [4.4, 7.0, 2.0],
              [8.9, 1.4, 3.2]])

I could simply loop through the indices but I was hoping that there were a more efficient approach:

m = 3
out = np.empty((len(I), m))
for j, i in enumerate(I):
    out[j] = T[j : j+m]

Use `s = T; l = m; i = I` with the linked Q&A's answer. – Divakar Aug 29 '20 at 15:22 — Divakar, Aug 29 '20 at 15:22

score 0 · Answer 1 · answered Aug 31 '20 at 01:48

Here's an answer that does not require installing additional dependencies:

def rolling_window(a, window):
    a = np.asarray(a)
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)

    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)

m = 3
rolling_T = rolling_window(T, m)  # Note that this is a very efficient view
out = rolling_T[I]

numpy select multiple ranges from a 1d array

1 Answers1