Assigning to slices of 2D NumPy array

Question

I want to assign 0 to different length slices of a 2d array.

Example:

import numpy as np    

arr = np.array([[1,2,3,4],
                [1,2,3,4],
                [1,2,3,4],
                [1,2,3,4]])

idxs = np.array([0,1,2,0])

Given the above array arr and indices idxs how can you assign to different length slices. Such that the result is:

arr = np.array([[0,2,3,4],
                [0,0,3,4],
                [0,0,0,4],
                [0,2,3,4]])

These don't work

slices = np.array([np.arange(i) for i in idxs])
arr[slices] = 0

arr[:, :idxs] = 0

cs95 · Accepted Answer · 2018-02-20 00:26:32Z

5

You can use broadcasted comparison to generate a mask, and index into arr accordingly:

arr[np.arange(arr.shape[1]) <= idxs[:, None]] = 0

print(arr) 
array([[0, 2, 3, 4],
       [0, 0, 3, 4],
       [0, 0, 0, 4],
       [0, 2, 3, 4]])

answered Feb 20, 2018 at 0:26

cs95

406k106 gold badges744 silver badges797 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

James Schinner Over a year ago

That's awesome! Once I see the solution it looks simple, yet I could not figure it out.

Ferdinando Randisi Over a year ago

Cool stuff! But why do you need the None within the brackets?

cs95 Over a year ago

@FerdinandoRandisi That's how I specify that I wish to broadcast the operation across multiple rows :)

Ferdinando Randisi · Accepted Answer · 2018-02-20 00:31:15Z

1

This does the trick:

import numpy as np    

arr = np.array([[1,2,3,4],
               [1,2,3,4],
               [1,2,3,4],
               [1,2,3,4]])
idxs = [0,1,2,0]
for i,j in zip(range(arr.shape[0]),idxs):
  arr[i,:j+1]=0

answered Feb 20, 2018 at 0:31

Ferdinando Randisi

4,4177 gold badges36 silver badges46 bronze badges

Comments

berna1111 · Accepted Answer · 2018-02-20 00:50:12Z

1

import numpy as np

arr = np.array([[1, 2, 3, 4],
                [1, 2, 3, 4],
                [1, 2, 3, 4],
                [1, 2, 3, 4]])

idxs = np.array([0, 1, 2, 0])

for i, idx in enumerate(idxs):
    arr[i,:idx+1] = 0

answered Feb 20, 2018 at 0:50

berna1111

1,8611 gold badge18 silver badges24 bronze badges

Comments

Paul Panzer · Accepted Answer · 2018-02-20 01:04:44Z

1

Here is a sparse solution that may be useful in cases where only a small fraction of places should be zeroed out:

>>> idx = idxs+1
>>> I = idx.cumsum()
>>> cidx = np.ones((I[-1],), int)
>>> cidx[0] = 0
>>> cidx[I[:-1]]-=idx[:-1]
>>> cidx=np.cumsum(cidx)
>>> ridx = np.repeat(np.arange(idx.size), idx)
>>> arr[ridx, cidx]=0
>>> arr
array([[0, 2, 3, 4],
       [0, 0, 3, 4],
       [0, 0, 0, 4],
       [0, 2, 3, 4]])

Explanation: We need to construct the coordinates of the positions we want to put zeros in.

The row indices are easy: we just need to go from 0 to 3 repeating each number to fill the corresponding slice.

The column indices start at zero and most of the time are incremented by 1. So to construct them we use cumsum on mostly ones. Only at the start of each new row we have to reset. We do that by subtracting the length of the corresponding slice such as to cancel the ones we have summed in that row.

edited Feb 20, 2018 at 1:04

answered Feb 20, 2018 at 0:47

Paul Panzer

53.3k3 gold badges60 silver badges103 bronze badges

2 Comments

Paul Panzer Over a year ago

@JamesSchinner I've added a bit of explanation.

James Schinner Over a year ago

Much appreciated. This really does seem like the way to go for large arrays.

Collectives™ on Stack Overflow

Assigning to slices of 2D NumPy array

4 Answers 4

3 Comments

Comments

Comments

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

3 Comments

Comments

Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related