Search code examples
pythonnumpy-slicing

How to loop over subarrays until the end of a numpy array

Suppose that we have a numpy array:

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

What I'm looking for is a way to loop over subranges in the array starting from a position from the end. For example like this:

[4 5]
[6 7]
[8 9]

The following code does not work:

n = 6
m = 2
for i in range(0, n, m):
    print(a[-n+i:-n+i+m])

It prints:

[4 5]
[6 7]
[]

The problem is that in the last iteration the value to the right of : becomes 0. I can handle this with a condition like:

if -n+i+m == 0:
    print(a[-n+i:])

But I'm looking for a more elegant solution if possible.

Solution

  • Get slice starting from the needed position and reshape:

    a[-n:].reshape(-1, m)

    array([[4, 5],
       [6, 7],
       [8, 9]])