Numpy数组按最后一列分组的列求和实现（无循环）

嘿，这事儿用Numpy的原生工具就能轻松搞定，完全不用写for循环绕圈子。咱们直接上代码，再一步步拆解逻辑：

import numpy as np
arr = np.array([[1,2,3,4,2000], [5,6,7,8,2000], [9,0,1,2,2001], [3,4,5,6,2001], [7,8,9,0,2002], [1,2,3,4,2002], [5,6,7,8,2003], [9,0,1,2,2003] ])

# 1. 提取分组依据：数组的最后一列
groups = arr[:, -1]
# 2. 获取唯一分组键，以及每个分组在原数组中首次出现的索引
unique_groups, group_indices = np.unique(groups, return_index=True)
# 3. 按索引对数组进行分段求和
grouped_sum = np.add.reduceat(arr, group_indices, axis=0)

print(grouped_sum)

运行这段代码，你会得到想要的结果：

array([[ 6,  8, 10, 12, 4000],
       [12,  4,  6,  8, 4002],
       [ 8, 10, 12,  4, 4004],
       [14,  6,  8, 10, 4006]])

逻辑拆解： #

• 提取分组键：arr[:, -1]直接取所有行的最后一列，这就是我们的分组标识。
• 定位分组起始点：np.unique(..., return_index=True)会返回唯一的分组值，以及每个值第一次出现在原数组中的索引。因为你的原数组已经是按分组键连续排列的，所以这些索引正好是每个分组的起始位置（比如2000对应索引0，2001对应索引2，以此类推）。
• 分段求和：np.add.reduceat是Numpy专门用来分段聚合的函数，这里指定axis=0表示按行方向，根据group_indices的位置把数组切成段，然后每段内部求和，完美实现分组求和的需求。

如果你的原始数组不是按分组键连续排列的，也没关系——只需要先按最后一列排序，再用上面的方法就行，比如加一步arr = arr[arr[:, -1].argsort()]，剩下的逻辑完全一样。

内容的提问来源于stack exchange，提问作者Infinity Cliff