How to set the values of different channels in a numpy array to zero based on index values from anot-CodePudding

I have a numpy array "arr" and an array of indices "ind":


import numpy as np

arr = np.random.randint(255, size=(100,64,64,16))

ind = np.random.randint(16, size=(100,2))

In arr, the last dimension represents channels while the first dimension represents number of samples. In ind, each row represents the two random channel indices corresponding to every sample in arr. I wish to make the values of entire 64 x 64 channel equal to 0 for every sample in arr corresponding to channel index in ind, without using any loop. How can this be achieved?

I have tried using:

arr[:,:,:,ind] = 0

I thought that indices would be broadcasted as per the sample, but instead entire array becomes 0. Using loops is quite time consuming and inefficient. I also wanted to use np.where, but I am not sure what condition to use to access the indices of arrays.

CodePudding user response：

I believe that you can use np.put_along_axis for this:

import numpy as np

arr = np.random.randint(255, size=(100,64,64,16))
ind = np.random.randint(16, size=(100,2))

np.put_along_axis(arr, ind[:, None, None, :], 0, axis=3)