python - 基于 Numpy 索引的行总和

Question

我有一个由二维向量和一维索引数组组成的二维数组。

如何使用 numpy 添加/求和共享相同索引的二维数组的行？

例子:

arr = np.array([[48, -51], [-15, -55], [26, -49], [-13, -17], [-67,  -7], [23, -48], [-29, -64], [37,  68]])
idx = np.array([0, 1, 1, 2, 2, 3, 3, 4])

#desired output
array([[48, -51],
       [11, -104],
       [-80, -24],
       [-6, -112],
       [ 37, 68]])

注意原始数组 arr 的形状是 (8, 2)，运算结果是 (5, 2)。

Answer 1

如果索引并不总是分组，则应用 np.argsort第一:

order = np.argsort(idx)

您可以使用 np.diff 计算总和的位置其次是 np.flatnonzero得到指数。我们还将在前面加上零并将所有内容移动 1:

breaks = np.flatnonzero(np.concatenate(([1], np.diff(idx[order])))

breaks 现在可以用作 np.add.reduceat 的参数:

result = np.add.reduceat(arr[order, :], breaks, axis=0)

如果索引已经分组，则根本不需要使用order:

breaks = np.flatnonzero(np.concatenate(([1], np.diff(idx)))
result = np.add.reduceat(arr, breaks, axis=0)