How to calculate conditional sum of a part of vector/array based on a 2nd one

Hello,

I want to get the conditional sum of part of a vector based on a second vector.

For example:

A = [143 14 -4 299 -29 83 151 190 178 -61 -109 ]

and

B = [1 1 -1 1 -1 1 1 1 1 -1 -1 ]

then I want to get the output:

C=[0 157 -4 299 -29 0 0 0 602 -170 ]

and

D=[0 2 1 1 1 0 0 0 4 0 2 ]

The idea is, when successive B's = 1, I add the corresponding values from A and store in C. Similarly if successive B's = -1, I add the values up and store in C. If more than 1 value in A is added, I fill the gaps in C with 0's. It would also be great if I could get a seperate array D mentioning how many values in A have been added to produce the value in C!

Please help! Thanks a lot!

[Cedric]

How large are these vectors? There is a vector way to do it (requires quite a few operations) and a way based on a simple FOR loop. The vector way is likely to be very efficiency on large vectors/arrays, and the FOR loop is likely to be more efficient on small vectors/arrays.

Answer 1

[Teja Muppirala]

Just for comparison, here is a FOR loop solution. FOR loops can actually be very fast, even for large arrays, especially when you are doing simple operations like this.

C = zeros(size(A));
D = zeros(size(A));

previous = B(1);
runningSum = A(1);
numValues = 1;
for n = 2:numel(B)
    if B(n) == previous
        runningSum = runningSum + A(n);
        numValues = numValues+1;
    else
        C(n-1) = runningSum;
        D(n-1) = numValues;
        runningSum = A(n);
        numValues = 1;
        previous = B(n);
    end
end
C(end) = runningSum;
D(end) = numValues;

Answer 2

[Cedric]

The following is an example of "vector approach":

 dif         = [true, diff(B)~=0] ;
 blockId     = cumsum(dif) ;
 blockStart  = find(dif) ;
 blockEnd    = [blockStart(2:end)-1, numel(A)] ;

 blockSum    = accumarray(blockId(:), A(:)) ;
 C           = zeros(size(A)) ;
 C(blockEnd) = blockSum ;

 blockSize   = accumarray(blockId(:), ones(size(A))) ;
 D           = zeros(size(A)) ;
 D(blockEnd) = blockSize ;

Running this produces:

 >> C
 C =
     0   157    -4   299   -29     0     0     0   602     0  -170

 >> D
 D =
     0     2     1     1     1     0     0     0     4     0     2

As you can see, a simpler FOR loop is likely to be more efficient for small array sizes, but this vector approach certainly wins if A and B are large.

I don't have time to check it extensively though, and I leave that to you. Check well boundary cases in particular.

Note that you can probably simplify the first part by computing directly blockEnd without defining blockStart, but I leave it this way for sake of clarity (or more honestly by lack of time ;-)).