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!