如何为NumPy数组添加分组滚动计数列（替代Pandas实现）

需求说明 #

我有一个包含4列的DataFrame，想要新增一列计算分组内的滚动累计值。虽然知道可以用Pandas的groupby结合cumcount() + 1实现，但希望把数据转入NumPy数组，通过向量化操作完成这个需求。

示例数据 #

• 输入数据：包含用户ID、日期等4列，按用户ID分组，每组内有若干条按日期排序的记录
• 输出数据：新增一列（比如“投诉总数”），每组内从第一条到当前行的累计计数——同一用户的第1条记录值为1，第2条为2，以此类推

实现方案 #

步骤1：从Pandas转NumPy数据 #

假设DataFrame的分组键是user_id列，先将分组键和数据转为NumPy数组：

import numpy as np
import pandas as pd

# 模拟示例数据
df = pd.DataFrame({
    'user_id': [1,1,1,2,2,3],
    'date': ['2023-01-01','2023-01-02','2023-01-03','2023-01-01','2023-01-02','2023-01-01'],
    'col3': ['a','b','c','d','e','f'],
    'col4': [10,20,30,40,50,60]
})

# 提取分组键的NumPy数组
group_keys = df['user_id'].to_numpy()

步骤2：向量化计算分组累计值 #

核心通过排序、找分组边界、生成累计序列再还原顺序实现：

# 获取排序后的索引与分组键
sorted_idx = np.argsort(group_keys)
sorted_groups = group_keys[sorted_idx]

# 定位分组变化的边界
group_changes = np.concatenate([[True], sorted_groups[1:] != sorted_groups[:-1]])
group_starts = np.where(group_changes)[0]
# 计算每个分组的长度
group_lengths = np.diff(np.concatenate([group_starts, [len(sorted_groups)]]))

# 生成各分组内的累计计数并拼接
cumulative_counts = np.concatenate([np.arange(1, length+1) for length in group_lengths])

# 还原回原始数据的顺序
result = np.empty_like(cumulative_counts)
result[sorted_idx] = cumulative_counts

# 将结果添加回DataFrame
df['累计值'] = result

结果验证 #

运行后df新增的“累计值”列，和df.groupby('user_id').cumcount() + 1的结果完全一致，但全程采用NumPy向量化操作，避免了Pandas分组循环的额外开销。

内容的提问来源于stack exchange，提问作者Pawan Tolani