Convert ndarray to 2D ndarray

Given an ndarray:

import numpy as np
import random

idxs = np.arange(10)
np.random.shuffle(idxs)

I get an output:

array([1, 2, 7, 8, 5, 9, 3, 4, 6, 0])

I would like to convert this to an array of ndarrays that looks like:

array([[1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [7, 8, 9, 10, 11],
       [8, 9, 10, 11, 12],
       ...])

where for each element of the first array, I want the indexes up to 5 above it. How can I do this? I have tried using

np.reshape(idxs, (idxs.size, 1))

to get it into this:

array([[1],
       [2],
       [7],
       [8],
       ...])

but then I would like to get all indexes up to 5 above it, as mentioned above.

I hope I explained this clearly, grateful for any help.

Use broadcasting:

idxs = np.array([1, 2, 7, 8, 5, 9, 3, 4, 6, 0])

out = idxs[:,None] + np.arange(5)

# or
out = idxs.reshape((-1, 1))+np.arange(5)

Alternatively (mostly for fun as this is not as efficient), if your indices are 0-n:

from numpy.lib.stride_tricks import sliding_window_view as swv

out = swv(np.arange(10+5), 5)[idxs]

(If not, this would generate many intermediate rows unnecessarily)

Output:

array([[ 1,  2,  3,  4,  5],
       [ 2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11],
       [ 8,  9, 10, 11, 12],
       [ 5,  6,  7,  8,  9],
       [ 9, 10, 11, 12, 13],
       [ 3,  4,  5,  6,  7],
       [ 4,  5,  6,  7,  8],
       [ 6,  7,  8,  9, 10],
       [ 0,  1,  2,  3,  4]])

I have this simple way of getting a (n, 1) shaped nd-array out of your code by using the np.newaxis.

In [3]: np.arange(10)[:, np.newaxis]
Out[3]:
array([[0],
       [1],
       [2],
       [3],
       [4],
       [5],
       [6],
       [7],
       [8],
       [9]])