CODEWITHSUNDEEP
pythonarraysnumpyfor-loopbroadcast
andro_analytics
andro_analytics

Reputation: 23

Numpy way of step forward for loop

I want to avoid for-loop with step in the following code and replace it with Numpy code to speed up the process:

import numpy as np

A = np.array([[0,0,7,0,0,0], [0,0,0,0,5,0]])

for i in range(A.shape[0]):
    for j in range(A.shape[1]-1):
        if A[i,j]==7 and A[i,j+1]==0:
            A[i,j+1]=7

I know how to do it with for-loop without step. Say, A,B,C are 2D arrays with same size, then this slow code:

for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        if A[i,j]==7 and B[i,j]==0:
            C[i,j]=7

...can be faster by the following single line numpy code:

C[(A==7) & (B==0)]=7

I guess it should be something similar, including np.where and np.roll functions? Appreciate your help!

Upvotes: 1

Views: 149

Answers (1)

Ehsan
Ehsan

Reputation: 12407

I am pretty sure there are better ways, but in case you did not find it, here is a faster way:

for i in range(A.shape[1]-1):
  A = np.where((A==0)&(np.pad(A,((0,0),(1,0)))[:,:-1]==7),7,A)

output:

[[0 0 7 7 7 7]
 [0 0 0 0 5 0]]

Upvotes: 1

Related Questions