Modify different columns in each row of a 2D NumPy array

Question

I have the following problem:

Let's say I have an array defined like this:

A = np.array([[1,2,3],[4,5,6],[7,8,9]])

What I would like to do is to make use of Numpy multiple indexing and set several elements to 0. To do that I'm creating a vector:

indices_to_remove = [1, 2, 0]

What I want it to mean is the following:

1. Remove element with index '1' from the first row
2. Remove element with index '2' from the second row
3. Remove element with index '0' from the third row

The result should be the array [[1,0,3],[4,5,0],[0,8,9]]

I've managed to get values of the elements I would like to modify by following code:

values = np.diagonal(np.take(A, indices, axis=1))

However, that doesn't allow me to modify them. How could this be solved?

Answer 1

You could use integer array indexing to assign those zeros -

A[np.arange(len(indices_to_remove)), indices_to_remove] = 0

Sample run -

In [445]: A
Out[445]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

In [446]: indices_to_remove
Out[446]: [1, 2, 0]

In [447]: A[np.arange(len(indices_to_remove)), indices_to_remove] = 0

In [448]: A
Out[448]: 
array([[1, 0, 3],
       [4, 5, 0],
       [0, 8, 9]])