cumsum restarting at 0 column-wise

I have a large matrix with all zeros and ones, and I want to calculate the cumsum (down the columns), but I want the sum to restart when there is a 0. Example:

INPUT = [0 0 1 0 0 1; 1 1 1 1 1 1; 1 1 1 1 1 1; 0 1 1 0 1 0; 1 1 1 1 1 1; 0 0 1 0 0 1]

OUTPUT = [0 0 1 0 0 1; 1 1 2 1 1 2; 2 2 3 2 2 3; 0 3 4 0 3 0; 1 4 5 1 4 1; 0 0 6 0 0 2]

I have seen some code to do this on one vector, but not on all columns of a matrix at the same time. Obviously, it's trivial to do by looping over the columns, but this is to be avoided for reasons of speed.

Any help is appreciated!

[Rik]

I would imagine the FINDSEQ function on the FEX could be adapted to this.

Answer 1

[Stephen23]

Download and use Matt Fig's excellent rcumsum, which is described as "cumulative sum of elements, restarted after every zero". It is used just like cumsum:

>> M = [0 0 1 0 0 1; 1 1 1 1 1 1; 1 1 1 1 1 1; 0 1 1 0 1 0; 1 1 1 1 1 1; 0 0 1 0 0 1];
>> rcumsum(M,1)
ans =
   0   0   1   0   0   1
   1   1   2   1   1   2
   2   2   3   2   2   3
   0   3   4   0   3   0
   1   4   5   1   4   1
   0   0   6   0   0   2

Answer 2

[Jan]

Another short solution:

In = [0 0 1 0 0 1; 1 1 1 1 1 1; 1 1 1 1 1 1; 0 1 1 0 1 0; 1 1 1 1 1 1; 0 0 1 0 0 1]

Out = In;
for k = 2:size(In, 1);
  Out(k, :) = (Out(k - 1, :) + Out(k, :)) .* Out(k, :);
end

For the general case, where the elements can differ from 0 and 1:

Out = In;
for k = 2:size(In, 1);
  Out(k, :) = (Out(k - 1, :) + Out(k, :)) .* (Out(k, :) ~= 0);
end