I have 2 numpy arrays:
a = np.array([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
b = np.array([2, 1, 2])
I want to use b
as starting indices into the columns of a
and set all the values of a
from those column indexes onwards to 0 like this:
np.array([[1, 2, 3],
[4, 0, 6],
[0, 0, 0]])
i.e., set elements of column 1 from position 2 onwards to 0, set elements of column 2 from position 1 onwards to 0, and set elements of column 3 from position 2 onwards to 0.
When I try this:
a[:, b:] = 0
I get
TypeError: only integer scalar arrays can be converted to a scalar index
Is there a way to slice using an array of indices without a for loop?
Edit: updated the example to show the indices can be arbitrary
You can use boolean array indexing. First, create a mask of indices you want to set to 0 and then apply the mask to array and assign the replacement value (e.g., 0 in your case).
mask = b>np.arange(a.shape[1])[:,None]
a[~mask]=0
output:
array([[1, 2, 3],
[4, 0, 6],
[0, 0, 0]])