CODEWITHSUNDEEP
pythonarraysnumpy
Stef1611
Stef1611

Reputation: 2387

Conditionnal Loop over a numpy array?

I am a beginner in Python and I usually program in C.

So, I have a numpy 2D array. I do the mean of the (i,j),(i+1,j),(i,j+1) and (i+1,j+1) values and I sum this mean if it is above a chosen value.

This is my python code :

Z=np.array([[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15]])
sum=0.
value=7.
for i in range(np.shape(Z)[0]-1):
    for j in range(np.shape(Z)[1]-1):
        a = (Z[i,j] + Z[i+1,j] + Z[i,j+1] + Z[i+1,j+1]) / 4
        if (a>=value):
            sum+=a
print (sum)

I know It does not sound very pythonic. How can I write it in pythonic way to speed up this code on a large 2D numpy array ?

Thanks for answer

Upvotes: 0

Views: 57

Answers (1)

John Zwinck
John Zwinck

Reputation: 249203

I'd do it this way:

quads = Z[:-1,:-1] + Z[1:,:-1] + Z[:-1,1:] + Z[1:,1:]
sum = quads[quads >= value * 4].sum() / 4

The first line computes the entire (x-1,y-1) array of sums of 2x2 elements:

array([[16, 20, 24, 28],
       [36, 40, 44, 48]])

The second line compares each of those 8 elements with value * 4, rather than dividing quads / 4 which would create another array of the same size unnecessarily. This lets us do a single scalar multiply and a scalar divide at the end, instead of an array divide. But you could also write it this way if you don't care about the optimization:

quads /= 4
sum = quads[quads >= value].sum()

Upvotes: 3

Related Questions