Numpy conditional but only in some indices?

Question

I have an array of shape [10,200,50]. I would like to replace all values which are:

1. Greater than 33
2. Fall within a set of indices on the third axis: indices=[0,1,15,20,19]

So - any value which has any of those indices on axis 3 AND is greater than 33 will be replaced by the number 22.

I know this is quite trivial, but I'm having trouble searching the right terms to find the solution. My instinct was to do: arr[arr[:,:,indices]==33]]=22, but this doesn't work because the shape of the internal array does not match the shape of the outer.

Answer 1

You can create a mask and set true where the condition is met.

import numpy as np
data = np.random.randn(10, 200, 50)
threshold = 33
new_value = 34
indices = [0,1,15,20,19]
data[0, 0, 0] = 100
mask = np.zeros(data.shape, dtype=bool)
mask[:, :, indices] = data[:, :, indices] > threshold
data[mask] = new_value
print(data[0, 0 ,0])