Suppose I have two numpy arrays a and b of shape (n, ) and a boolean array c, also of shape (n, ).
I want to compute a shape (n, ) numpy array x, where x[i] = a[i] if c[i], else b[i].
E.g.
a = np.array([3, 4, 5])
b = np.array([-1, -5, -9])
c = np.array([True, False, True])
x = np.array([3, -5, 5]).
Do anyone know how to do this with numpy operations?
Thank you!
Using numpy.where
>>> import numpy as np
>>> a = np.array([3, 4, 5])
>>> b = np.array([-1, -5, -9])
>>> c = np.array([True, False, True])
>>>
>>> np.where(c, a, b)
array([ 3, -5, 5])
Using element-wise multiplication
>>> x = np.array([3, 4, 5])
>>> y = np.array([-1, -5, -9])
>>> c = np.array([True, False, True])
>>> x * c + y * (1 - c)
array([ 3, -5, 5])
np.where.You can use the zip method towork on the three lists at the same time in your list comprehension:
[aa if cc else bb for aa, bb, cc in zip(a, b, c)]
