Replace 2D array elements with zeros, using a column index vector

Question

I have a 2D array A:

28  39  52
77  80  66   
 7  18  24    
 9  97  68

And a vector array of column indexes B:

1   
0   
2    
0

How in a pythonian way, using base python or numpy, can I replace with zeros, elements on each row of A that correspond to the column indexes in B?

I should get the following new array A, with one element on each row (the one in the column index stored in B), replaced with zero:

28   0  52
 0  80  66   
 7  18   0    
 0  97  68

Thanks for your help!

Answer 1

Here's a simple, pythonian way using enumerate:

for i, j in enumerate(B):
    A[i][j] = 0

Answer 2

In [117]: A = np.array([[28,39,52],[77,80,66],[7,18,24],[9,97,68]])                                                     
In [118]: B = [1,0,2,0]                                                                                                 

To select one element from each row, we need to index the rows with an array that matches B:

In [120]: A[np.arange(4),B]                                                                                             
Out[120]: array([39, 77, 24,  9])

And we can set the same elements with:

In [121]: A[np.arange(4),B] = 0                                                                                         
In [122]: A                                                                                                             
Out[122]: 
array([[28,  0, 52],
       [ 0, 80, 66],
       [ 7, 18,  0],
       [ 0, 97, 68]])

This ends up indexing the points with indices (0,1), (1,0), (2,2), (3,0).

A list based 'transpose' generates the same pairs:

In [123]: list(zip(range(4),B))                                                                                         
Out[123]: [(0, 1), (1, 0), (2, 2), (3, 0)]