I have a multidimensional NumPy array (16x212), and I want to calculate the cumulative sum for each of the 212 columns - however, the calculation should restart at 0 if there was a 0 in between.
e.g. array([0, 1, -1, 1, 0, -1, 1, 0, 0, 1, 0, -1, 0, 0, 1, -1], ...)
I watched the following video where you cumulated an array of 0's and 1's together, however the method doesn't work any more if negative numbers are present.
Can someone help me please?
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.
