Equation of vectors with booleans in python

Question

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?

Answer 1

with numpy :

import numpy as np
a = np.array([0,1,2,3,4,5])
b = np.array([1,1,2,2,3,3])
a-b+5*((a==0)|(b==0))
#> array([4, 0, 0, 1, 1, 2])

Answer 2

Lots and lots of ways. Here's one:

a = range(6)
b = [1, 1, 2, 2, 3, 3]
r = [ai - bi + 5*(ai==0 or bi==0) for ai, bi in zip(a,b)]

Answer 3

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))