CODEWITHSUNDEEP
pythonnumpy
flowfree
flowfree

Reputation: 16462

How to iterate Numpy array and perform calculation only if element matches a criteria?

I want to iterate a numpy array and process only elements match with specific criteria. In the code below, I want to perform calculation only if element is greater than 1.

a = np.array([[1,3,5],
              [2,4,3],
              [1,2,0]])

for i in range(0, a.shape[0]):
    for j in range(0, a.shape[1]):
        if a[i,j] > 1:
            a[i,j] = (a[i,j] - 3) * 5

Is it possible to use single-line code instead of the double loop above? and perhaps make it faster?

Upvotes: 0

Views: 324

Answers (2)

donghyun208
donghyun208

Reputation: 4857

for index, x in np.ndenumerate(a):
    if x > 1:
        a[index] = (a[index] - 3) * 5

Upvotes: 0

DSM
DSM

Reputation: 353329

Method #1: use a boolean array to index:

>>> a = np.array([[1,3,5], [2,4,3], [1,2,0]])
>>> a[a > 1]  = (a[a > 1] - 3) * 5
>>> a
array([[ 1,  0, 10],
       [-5,  5,  0],
       [ 1, -5,  0]])

This computes a > 1 twice, although you could assign it to a variable instead. (In practice it's very unlikely to be a bottleneck, of course, although if a is large enough memory can be an issue.)

Method #2: use np.where:

>>> a = np.array([[1,3,5], [2,4,3], [1,2,0]])
>>> np.where(a > 1, (a-3)*5, a)
array([[ 1,  0, 10],
       [-5,  5,  0],
       [ 1, -5,  0]])

This only computes a > 1 once, but OTOH computes (ax-3)*5 for every element ax in a, instead of only doing it for those elements that really need it.

Upvotes: 3

Related Questions