CODEWITHSUNDEEP
pythonarraysnumpyindexing
bmiselis
bmiselis

Reputation: 412

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

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?

Upvotes: 3

Views: 1750

Answers (1)

Divakar
Divakar

Reputation: 221534

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]])

Upvotes: 2

Related Questions