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

user_1389840 0
Viewed 652

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

matrix-indexing slice numpy arrays python
2 Answers

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

I think the issue is in a[:,b:]; here b: means little if b is not a scaler e.g. 5: means 6th onwards but [1,2,3]: means nothing when array is 2d.

It should be a[:,b]. Setting a[:,b] = 0 will set all columns specified in b to 0. Following is the run.


In [2]: import numpy as np

In [3]: a = np.array([[1, 2, 3],
   ...:               [4, 5, 6],
   ...:               [7, 8, 9]])
   ...: 
   ...: b = np.array([2, 1, 2])
   ...: 

In [4]: a
Out[4]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

In [5]: b
Out[5]: array([2, 1, 2])

In [6]: b.dtype
Out[6]: dtype('int64')

In [7]: a[:, b:] = 0
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-7-6e5050513225> in <module>
----> 1 a[:, b:] = 0

TypeError: only integer scalar arrays can be converted to a scalar index

In [8]: a[:, b] = 0

In [9]: a
Out[9]: 
array([[1, 0, 0],
       [4, 0, 0],
       [7, 0, 0]])

But that's not what you want.

To get what you want, you need to specify rows indices and column indices e.g. (1,1), (2,0), (2,1), (2,2).

In [11]: a[[1,2,2,2], [1, 0, 1,2]] = 0

In [12]: a
Out[12]: 
array([[1, 2, 3],
       [4, 0, 6],
       [0, 0, 0]])
user_18051290
Related