pr0crastin8r
pr0crastin8r

Reputation: 391

How do I do conditional array arithmetic on a numpy array?

I'm trying to get a better grip on numpy arrays, so I have a sample question to ask about them:

Say I have a numpy array called a. I want to perform an operation on a that increments all the values inside it that are less than 0 and leaves the rest alone. for example, if I had:

a = np.array([1,2,3,-1,-2,-3])

I would want to return:

([1,2,3,0,-1,-2])

What is the most compact syntax for this?

Thanks!

Upvotes: 16

Views: 30000

Answers (2)

John Kugelman
John Kugelman

Reputation: 361849

To mutate it:

a[a<0] += 1

To leave the original array alone:

a + (a<0)

Upvotes: 12

unutbu
unutbu

Reputation: 880259

In [45]: a = np.array([1,2,3,-1,-2,-3])

In [46]: a[a<0]+=1

In [47]: a
Out[47]: array([ 1,  2,  3,  0, -1, -2])

Upvotes: 41

Related Questions