Change different columns in each row of 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 use Numpy multiple indexing and set multiple elements to 0. To do this, I create a vector:
indices_to_remove = [1, 2, 0]
I want it to mean the following:
- Remove item at index '1' from first row
- Remove item with index "2" from second line
- Remove element at index '0' from third row
The result should be an array [[1,0,3],[4,5,0],[0,8,9]]
I was able to get the values of the elements that I would like to change using the following code:
values = np.diagonal(np.take(A, indices, axis=1))
However, this prevents me from changing them. How can this be solved?
+3
source to share
1 answer
You can use integer array indexing
to assign these zeros -
A[np.arange(len(indices_to_remove)), indices_to_remove] = 0
Example 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]])
+2
source to share