python - （Numpy 或 PyTorch）对给定的 bin 求和数组元素

Question

我希望使用 PyTorch 张量来解决这个问题。如果 torch 中没有有效的解决方案，请随意提出一个 numpy 解决方案。

让a是一维张量（或 numpy 数组），并且是 0 和排除bin_indices之间的整数的张量（np 数组） 。n我想计算bins位置处的数组i包含元素的总和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

如果您还可以提供一种批量处理这项工作的方法，我将非常感谢您！

标签： pythonarraysnumpypytorch

Answer 1

不确定这是否是最好的方法，但这是另一种解决方案：

>>> bins = torch.unique(bins_indices)
>>> vfunc = np.vectorize( lambda x: torch.sum( a[ bins_indices == x ] ) )
>>> vfunc( bins )
array([10,  3, -2])