pythonnumpy

Forward fill Numpy matrix / mask with values based on condition


I have the following matrix

import numpy as np


A = np.array([
    [0, 0, 0, 0, 1, 0, 1],
    [0, 0, 0, 0, 0, 0, 1],
    [1, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0]
]).astype(bool)

How do I fill all the rows column-wise after a column is True?

My desired output:

    [0, 0, 0, 0, 1, 1, 1],
    [0, 0, 0, 0, 0, 0, 1],
    [1, 1, 1, 1, 1, 1, 1],
    [0, 0, 0, 0, 0, 0, 0]

Solution

  • You could use logical_or combined with accumulate:

    np.logical_or.accumulate(A, axis=1)
    

    Output:

    array([[False, False, False, False,  True,  True,  True],
           [False, False, False, False, False, False,  True],
           [ True,  True,  True,  True,  True,  True,  True],
           [False, False, False, False, False, False, False]])
    

    If you want integers, go with maximum:

    np.maximum.accumulate(A.astype(int), axis=1)
    
    array([[0, 0, 0, 0, 1, 1, 1],
           [0, 0, 0, 0, 0, 0, 1],
           [1, 1, 1, 1, 1, 1, 1],
           [0, 0, 0, 0, 0, 0, 0]])