CODEWITHSUNDEEP
pythonarrayspython-3.xnumpyiteration
Jsevillamol
Jsevillamol

Reputation: 2543

Iterating over the last dimensions of a numpy array

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?

Upvotes: 0

Views: 251

Answers (1)

Joe Iddon
Joe Iddon

Reputation: 20414

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.

Upvotes: 1

Related Questions