How to replace values in a numpy array based on 2 other columns

Viewed 380

Let say I have the following array:

import numpy as np

a = np.random.randint(10, size=(3,5))

print(a)

Output looks like below:

[[4 0 3 7 3]
[6 5 7 3 3]
[6 6 4 2 1]]

If I want to select the row where column 0 element is 4 and replace column 2 element with 6, I can do the following.

a[a[:, 0] == 4, 2] = 6

Output will be as below:

[[4 0 6 7 3]
[6 5 7 3 3]
[6 6 4 2 1]]

How can I select the row where column 0 element is 6 AND column 1 element is 5 and replace column 2 element with 9 so that the output is as below:

[[4 0 6 7 3]
[6 5 9 3 3]
[6 6 4 2 1]]

2 Answers

You can do a bitwise AND on the conditions:

a[(a[:, 0] == 6) & (a[:, 1] == 5), 2] = 9

You can compare the array row-wise with a tuple/list and use all:

a[(a[:,:2]==(6,5)).all(axis=1), 2] = 9
Related