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

Question

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

I don't understand why `a[:, b:] = 0` and not `a[:, b] = 0` (all rows and columns specified by b)? — Dilawar, Jan 29 '21 at 00:42
Do you understand the error. In `[x:]` (a slice)`x` can only be an integer, not an array. — hpaulj, Jan 29 '21 at 00:49

score 3 · Accepted Answer · edited Jan 29 '21 at 02:02

3

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

edited Jan 29 '21 at 02:02

kmario23

57,311
13
161
150

answered Jan 29 '21 at 00:48

Ehsan

12,072
2
20
33

1

good use case of broadcasting. Alternatively, we can generate the exact boolean mask with `bmask = b<=np.arange(a.shape[1])[:,None]`, so we can avoid the mask negation. One less operation in terms of compute. – kmario23 Jan 29 '21 at 01:46
@kmario23 yeah, I thought about it first but then thought this way is more educating. Thank you for the comment and edit. – Ehsan Jan 29 '21 at 03:19

Dilawar · Answer 2 · 2021-01-29T00:50:35.120

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

Sorry, my example was too simple. I updated it to show the indices can be arbitrary. — user1389840, Jan 29 '21 at 00:35

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

2 Answers2