gustavgans
gustavgans

Reputation: 5191

Creating new array by comparing arrays

Can I create a new array by comparing arrays with numpy?

I have 3 arrays (A1, A2, A3). How can I find all indexes where A1 == 2 and A2 > A3 and write there the value 5 in a new array?

I have this matlab code which is doing this:

index = find(A1==2 & A2>A3);
new_array(index) = 5;

I found putmask and logical_and but not sure if this are the right tools and how do use it in my case. Thanks!

Upvotes: 0

Views: 69

Answers (2)

Manuel GOACOLOU
Manuel GOACOLOU

Reputation: 11

You can use the np.where function

import numpy as np

A1 = np.array([0, 0, 2, 2, 4, 4])
A2 = np.arange(len(A1))
A3 = np.ones(len(A1))*3

out = np.where((A1 == 2) & (A2 >= A3), 5, 0)

or more simple

out = ((A1 == 2) & (A2 >= A3))*5

Upvotes: 1

Ffisegydd
Ffisegydd

Reputation: 53698

The code below uses & to chain together conditionals. As such whenever A1 == 2 and A2 > A3 are both True, the index array will be True

import numpy as np

A1 = np.array([0, 0, 2, 2, 4, 4])
A2 = np.arange(len(A1))
A3 = np.ones(len(A1))*3

new_array = np.zeros(len(A1))

index = (A1 == 2) & (A2 > A3)

new_array[index] = 5
# array([ 0.,  0.,  5.,  5.,  0.,  0.])

You could use np.logical_and of course. But this restricts you to only two conditionals, whilst you can effectively chain as many as you like when using &.

Upvotes: 2

Related Questions