Summing with referenced index

I have 3 vectors - a sum vector, a contribution vector, and a value vector. I want to sum the value vectors according to their contribution vector and place them in their corresponding index in the sum vector. An example is:

A = [0;0] (sum vector), B = [0,0,1,1] (contribution vector) C=[20,30,40,10] (value vector)

Output: A = [20+30;40+10]

Such that the B vector is the same length as C and their corresponding index tell us what position in A to be added to.

I am able to achieve this by a for loop as such:

for index,value in enumerate(C):
    A[B[index]]+=value

However, as this will be part of my NN model forward loop it will cause significant performance issue. Specifically I was looking for a vector/matrix sorting approach that will be more efficient. In the example above, something that worked efficiently for me was:

A=torch.zeros(2,1)
C=C.reshape(2,2)
sum=torch.sum(C,1).reshape(2,1)
A += sum

However, I run into issues as it is not always the case that the indexes of A have the same contribution. For example - the case such that B = [0,0,0,1,1] and C=[20,30,40,10,50] . Is there a function or a strategic way to do this for general cases? Torch.gather seemed slightly similar but wasn’t able to do the job.Thanks!