CODEWITHSUNDEEP
pythonarraysnumpy
Berniyh
Berniyh

Reputation: 1

comparing numpy arrays element-wise setting an element-wise result

possibly this has been asked before, but I have a hard time finding a corresponding solution, since I can't find the right keywords to search for it.

One huge advantage of numpy arrays that I like is that they are transparent for many operations. However in my code I have a function that has a conditional statement in the form (minimal working example): import numpy as np

arr1 = np.array([1, 2, 3])
arr2 = np.array([1, 1, 3])

def func(x, y):
    if x > y:
        z = 1
    else:
        z = 2
    return z

func(arr1, arr2) obviously results in an error message: ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I do understand what the problem here is and why it can't work like this. What I would like to do here is that x > y is evaluated for each element and then an array z is returned with the corresponding result for each comparison. (Needs to ensure of course that the arrays are of equal length, but that's not a problem here) I know that I could do this by changing func such that it iterates over the elements, but since I'm trying to improve my numpy skills: is there a way to do this without explicit iteration?

Upvotes: 0

Views: 1314

Answers (1)

GPhilo
GPhilo

Reputation: 19123

arr1 > arr2 does exactly what you'd hope it does: compare each element of the two arrays and build an array with the result. The result can be used to index in any of the two arrays, should you need to. The equivalent of your function, however, can be done with np.where:

res = np.where(arr1 > arr2, 1, 2)

or equivalently (but slightly less efficiently), using the boolean array directly:

res = np.ones(arr1.shape)
res[arr1 <= arr2] = 2 # note that I have to invert the condition to select the cells where you want the 2

Upvotes: 2

Related Questions