Iterating over the last dimensions of a numpy array

Question

Consider the following snippet:

a = np.ones((2,2,2,2,2))
for example in a:
    for row in example:
        for col in row:
            for box in col:
                print (box.shape)

Having so many nested fors makes for some very ugly code.

How can I get the same effect with only one explicit iteration?

Joe Iddon · Accepted Answer · 2018-11-24T21:38:14.853

1

Reshape your array:

for box in a.reshape(16,2):
    print(box.shape)

A general solution:

for box in a.reshape(np.prod(a.shape[:-1]), a.shape[-1]):
     print(box.shape)

which will always iterate over one-before-the-last dimension of a.

edited Nov 24 '18 at 21:38

answered Nov 24 '18 at 21:20

Joe Iddon

20,101
7
33
54

Any way of doing this without referencing the shape of a, only the axis we want to iterate over? – Jsevillamol Nov 24 '18 at 21:34
1

@Jsevillamol See my update for a general solution :) – Joe Iddon Nov 24 '18 at 21:38

Iterating over the last dimensions of a numpy array

1 Answers1

Linked