[FIXED] (Numpy or PyTorch) Sum array elements for given bins

Issue

I would like this problem to be solved using PyTorch tensors. If there is no efficient solution in torch, then feel free to suggest a numpy solution.

Let a be a 1-dimensional tensor (or numpy array), and bin_indices be a tensor (np array) of integers between 0 and n excluded. I want to compute the array bins that at position i contains the sum of elements of a[bins_indices == i].

n = 3

a = [1, 4, 3, -2, 5]  # Values

bins_indices = [0, 0, 1, 2, 0]  # Correspondent bin indices

bins = [10, 3, -2]  # bins[0] = 1 + 4 + 5 etc. bins has 3 elements since n=3

If you can provide also a way of making this work for batches I would be immensely grateful to you!

Solution

Here's a one-line Numpy solution I could think of:

bins = [np.sum(a[np.argwhere(bins_indices == i).flatten()]) for i in range(len(a))]

Answered By - Stephen Mylabathula