如何用一行NumPy代码实现二维数组按首列分组求均值？

我有一个两列格式的文本文件，内容示例如下：

1 0.0169
1 0.0013
2 0.0070
2 0.0181
2 0.0008
2 0.0014
2 0.0005
4 0.0084
4 0.0154
5 0.0075
5 0.0171
5 0.0038
5 0.0115
5 0.0105
5 0.0094
6 0.0098
6 0.0129
7 0.0085
7 0.0097
7 0.0046
7 0.0021
7 0.0151
7 0.0117
...

通过G = np.loadtxt('2d.txt')将其读取为二维NumPy数组后，希望按第一列的值分组，将第二列替换为对应分组的平均值，例如：

• 第一列值为1的分组：平均值为(0.0169+0.0013)/2
• 第一列值为2的分组：平均值为(0.0070+0.0181+0.0008+0.0014+0.0005)/5
• 其余分组以此类推

需要用一行NumPy命令直接返回转换后的二维数组，而非通过循环等方式实现。

可以结合np.unique和np.bincount实现一行代码完成需求，以下是通用版本（支持第一列为任意可哈希类型的分组键）：

result = np.column_stack((uniq_vals := np.unique(G[:,0]), np.bincount(idx := np.unique(G[:,0], return_inverse=True)[1], weights=G[:,1]) / np.bincount(idx)))

代码拆解说明： #

1. np.unique(G[:,0], return_inverse=True)：提取第一列的唯一值，同时得到每个原元素对应的唯一值索引idx
2. np.bincount(idx, weights=G[:,1])：按索引分组，计算第二列的总和
3. np.bincount(idx)：计算每个分组的元素个数
4. 两者相除得到分组平均值，再通过np.column_stack将唯一值和对应平均值组合成二维数组

如果确定第一列都是非负整数，也可以用更简洁的版本：

result = np.column_stack((np.unique(G[:,0]), np.bincount(G[:,0].astype(int), weights=G[:,1]) / np.bincount(G[:,0].astype(int))))

内容的提问来源于stack exchange，提问作者Emil Frlez