elegant solution in numpy

Question

I have the following numpy array:

sample_size = 100000
z = np.random.randn(sample_size)

What would be an elegant way to say:

if z >= 0: f = 1 + z

if z < 0:  f = 1 / (2 - z)

user_4975981 · Answer

Use masking:

np.where(z>=0, 1+z, 1/(2-z))