I am trying to create an equation where you input an array or vector and it includes a boolean function. For the items in the array where this is satisfied (bool=True), then the equation proceeds to be solved in one way to produce another array.
I have attached here the similar code that works in R and want to do something similar in python
a <- c(0,1,2,3,4,5)
b <- c(1,1,2,2,3,3)
a-b+5*(a==0|b==0)
The output of that is a vector:
[1] 4 0 0 1 1 2
Does anyone know how to do something similar in python3, maybe with numpy?
>Solution :
A way using numpy:
import numpy as np
a = np.arange(6)
b = np.repeat(np.arange(1,4),2)
a-b+5*(np.equal(a,0) | np.equal(b,0))