CODEWITHSUNDEEP
pythonarraysloopsnumpy
AJaramillo
AJaramillo

Reputation: 440

Index order in loop with numpy array

Is there an specific order to make several loops in a numpy array in order to be fast? I mean, let be A a numpy array with A = A(i,j,k).

If I have to make loops in which order should I do the for sequence? Thanks

Upvotes: 1

Views: 1200

Answers (1)

hpaulj
hpaulj

Reputation: 231355

Are you looping over all dimensions, or just one?

for aa in A:
   print(aa)

iterates over the 1st dimension

e.g.

In [479]: A=np.arange(12).reshape(3,4)
In [480]: A
Out[480]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
In [481]: for a in A:
   .....:     print(a)
   .....:     
[0 1 2 3]
[4 5 6 7]
[ 8  9 10 11]

You can iterate by index over other dimensions

In [482]: for j in range(4):
   .....:     print(A[:,j])
   .....:     
[0 4 8]
[1 5 9]
[ 2  6 10]

or transpose:

In [483]: for a in A.T:
   .....:     print(a)
   .....:     
[0 4 8]
[1 5 9]
[ 2  6 10]
[ 3  7 11]

There's a minor speed advantage if you iterate over the first dimension, rather than the last. But I stress it is minor.

It best to avoid loops entirely, working with compiled methods that operate on the whole array at once. Iteration is, by comparison, much slower.

for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        for k in range(A.shape[2]):
             A[i,j,k]

will be horribly slow, regardless of the order of nesting the loops.

Upvotes: 1

Related Questions