pythonarraysnumpyslicematrix-indexing

Can I slice a numpy array using an array as indices?


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


Solution

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