user1482980
user1482980

Reputation: 167

Reducing loops in python

I would like to implement the following snippet using just the inntermost for loop(the one which iterates 3 times),as this consumes a lot of time.

for i in arange(r):
    for j in arange(c):    
        for k in arange(3):
            if m[i,j]==n[i,j,k]:
                 new[i,j]=old[i,j,k]

Could anyone suggest a better method?

Upvotes: 3

Views: 1159

Answers (3)

sega_sai
sega_sai

Reputation: 8538

for k in range(3):
    ind = m == n[:,:,k]
    new[ind] = old[:,:,k][ind]

Upvotes: 5

Emmanuel
Emmanuel

Reputation: 14209

Since arange(c) is computed for each i, and arange(3) for each couple (i, j), computing them once and for all outside the loop could save some time:

range_j = arange(c)
range_3 = arange(3)
for i in arange(r):
    for j in range_j:    
        for k in range_3:
            if m[i,j]==n[i,j,k]:
                 new[i,j]=old[i,j,k]

Of course this is only valid because these ranges are independent from i and j.

Upvotes: -1

Jon Clements
Jon Clements

Reputation: 142156

Look at using itertools.product - never used it with numpy arrays, but it might just work (and don't see why not)

for i, j, k in itertools.product(arange(r), arange(c), arange(3)):
    if m[i,j]==n[i,j,k]:
         new[i,j]=old[i,j,k]

Upvotes: 2

Related Questions