Map function with array based on integers on 2D numpy array-CodePudding

>>> a = 0.001 * np.random.randint(0, 1000, size=(5))
>>> a
array([0.524, 0.311, 0.603, 0.957, 0.923])

>>> b = np.random.randint(0, 2, size=(3, 5))
>>> b
array([[1, 1, 1, 0, 1],
       [0, 0, 1, 0, 0],
       [0, 0, 1, 0, 1]])

I want to map over each row of the 2D array b, and if the value at row[i] == 1, set it to a[i], if the value at row[i] == 0, I want to set it to 1 - a[i].

For example, for b[1] = row = [0, 0, 1, 0, 0], it should become:

[1 - 0.524, 1 - 0.311, 0.603, 1 - 0.957, 1 - 0.923]

i.e.

[0.476, 0.689, 0.603, 0.043, 0.077]

How can I achieve this for all rows of the 2D matrix without resorting to for loops? (i.e. leverage efficiency of numpy)

CodePudding user response：

There are a couple of ways of doing this. I would probably use the where and out arguments to np.subtract:

np.subtract(np.ones(len(b)), a, out=np.broadcast_to(a, b.shape).copy(), where=b.astype(bool))

Going with @hpaulj's solution to use the 3-arg version of np.where is probably much cleaner in this case:

np.where(b, a, 1 - a)