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

【问题标题】：(Numpy or PyTorch) Sum array elements for given bins（Numpy 或 PyTorch）对给定的 bin 求和数组元素
【发布时间】：2021-03-30 01:55:19
【问题描述】：

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

让a 是一维张量（或 numpy 数组），bin_indices 是 0 到 n 之间的整数张量（np 数组），排除在外。我想计算数组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

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

【问题讨论】：

标签： python arrays numpy pytorch

【解决方案1】：

这是我能想到的单行 Numpy 解决方案：

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

【讨论】：

【解决方案2】：

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

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

【讨论】：