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

Answer 1

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

Answer 2

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

Answer 3

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.

Answer 4

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