Modifying all elements on a row after a delimiter in a 2d numpy array

December 28, 2023

If I have a 2d numpy (integers and d is also an integer) array like

[[0 1 2 d]
 [3 4 d 5]
 [6 d 7 8]]

How can I 0 all elements (row-wise) after d (and including) on each row?

I’ve used a for loop for this but I was wondering if there was a vectorised method via numpy. I’ve just seen that maybe you could subtract a triangular matrix :

[[0 0 0 d]
 [0 0 d 5]
 [0 d 7 8]]

but this doesn’t solve my issue as it requires zeroing the values before d instead.

EDIT:

repr(array)

array([[0, 1, 2, d],
   [3, 4, d, 5],
   [6, d, 7, 8]], dtype=int64)

>Solution :

IIUC, you can craft a mask with cumsum and use it to mask the leading values with where:

d = 9
a = np.array([[0, 1, 2, d],
              [3, 4, d, 5],
              [6, d, 7, 8]])

out = np.where(np.cumsum(a == d, axis=1), a, 0)

Variant with cumprod and in place modification:

a[np.cumprod(a!=d, axis=1).astype(bool)] = 0

Output:

array([[0, 0, 0, 9],
       [0, 0, 9, 5],
       [0, 9, 7, 8]])

Intermediate masks (for the first non-zero values are considered True):

# np.cumsum(a == d, axis=1)
array([[0, 0, 0, 1],
       [0, 0, 1, 1],
       [0, 1, 1, 1]])

# np.cumprod(a!=d, axis=1).astype(bool)
array([[ True,  True,  True, False],
       [ True,  True, False, False],
       [ True, False, False, False]])