Matlab cumsum reset at NaN?

Question

If I have a vector of either 1's or NaN's like this:

[1 1 1 NaN 1 1 NaN 1 1 1 1]

How can I reset the cumsum to zero at the location of the NaNs like below:

[1 2 3 0 1 2 0 1 2 3 4]

Ideally I would like to have a vectorized solution since I need to do this for every column in a large matrix and the locations of the NaNs are not constant across the columns.

Thanks in advance.

Answer 1

You can try the following two 'vectorized' lines:

A(isnan(A)) = 1-diff([0 find(isnan(A))]);
cumsum(A)

ans =

  1     2     3     0     1     2     0     1     2     3     4

The trick is to substitute NaN with a value that will reset cumsum to 0 at those points.

Answer 2

I can only think of a few-pass solution:

v = [1 1 1 NaN 1 1 1 1 NaN 1];
a = v==v;              %% convert the values first to [1 1 1 0 1 1 1 1 0 1] format
n = a==0;              %% positions of the NaNs
c = cumsum(a);         %% your intermediate result
d = diff([0 c(n)]);    %% runs of ones
v(n) = -d;             %% replace Nans by -3, -4      [1 1 1 -3 1 1 1 1 -4 1] 
cumsum(v)              %% the answer [1 2 3 0 1 2 3 4 0 1]

Note: haven't checked extreme conditions (NaN in first/Last position, consecutive NaNs etc.)

Answer 3

There's a nice FEX file that treats this issue in this link. The code has user-specified treatment of NaNs. It allows the user to replace NaNs with zeros, or to skip over them, or to reset on NaNs, maintaining NaNs as placeholders.