Question

我有一个由2D向量和1D索引数组组成的2D数组。

如何使用numpy添加/总结2D数组共享相同索引的行？

示例：

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)

Answer 2

您可以将pandas用于此目的：

pd.DataFrame(arr).groupby(idx).sum().to_numpy()

输出：

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

基于Numpy的索引的行总数

2 个答案: