fill_zeros_with_last2d
is based on this answer.
The method here is to take the cumulated sum along the axes, then subtract the sum accumulated to the last zero from any later columns.
import numpy as np
def fill_zeros_with_last2d(adj, arr ):
row, col = np.indices( adj.shape )
col[(adj == 0) & (arr != 0) ] = 0
col = np.maximum.accumulate(col, axis = 1 )
# Find the last column with a value
return adj[row, col]
def cum_reset_2d( arr ):
accum = arr.cumsum( axis = 1 )
adj = accum * ( arr == 0 )
return accum - fill_zeros_with_last2d( adj, arr )
Testing on a small array
np.random.seed( 1235 )
arr = np.random.randint( -1, 2, size = ( 5, 10 ))
arr
"""
array([[ 1, 1, 1, 0, -1, -1, -1, 1, 1, -1],
[ 1, 1, -1, 1, 1, 1, -1, 0, -1, 1],
[ 0, -1, 1, 0, 0, -1, 1, 1, 0, -1],
[ 1, 0, -1, 0, -1, 1, 0, 0, 1, 1],
[-1, 0, 0, -1, -1, 1, 1, -1, 1, -1]])
"""
cum_reset_2d( arr )
"""
array([[ 1, 2, 3, 0, -1, -2, -3, -2, -1, -2],
[ 1, 2, 1, 2, 3, 4, 3, 0, -1, 0],
[ 0, -1, 0, 0, 0, -1, 0, 1, 0, -1],
[ 1, 0, -1, 0, -1, 0, 0, 0, 1, 2],
[-1, 0, 0, -1, -2, -1, 0, -1, 0, -1]])
"""
Note I had trouble getting this to work so it needs testing by before use. The original version didn't work for arr[3] but they all look OK now. There may still be some issues. I had adj
values of zero which had to be kept but were being overwritten.