Python: Insert columns into a numpy array based on mask

Suppose I have the following data:

mask = [[0, 1, 1, 0, 1]] # 2D mask
ip_array = [[4, 5, 2]
            [3, 2, 1]
            [1, 8, 6]] # 2D array

I want to insert columns of 0s into ip_array where ever there is 0 in the mask. So the output should be like:

[[0, 4, 5, 0, 2]
 [0, 3, 2, 0, 1]
 [0, 1, 8, 0, 6]]

I am new to numpy functions and I am looking for an efficient way to do this. Any help is appreciated!

>Solution :

Here’s one way to do it in two steps:

(i) Create an array of zeros of the correct shape (the first dimension of ip_array and the second dimension of mask)

(ii) Use the mask across the second dimension (as a boolean mask) and assign the values of ip_array to the array of zeros.

out = np.zeros((ip_array.shape[0], mask.shape[1])).astype(int)
out[..., mask[0].astype(bool)] = ip_array
print(out)

Output:

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