如何基于列行数限制与另一列总和阈值分组？Pandas问题求助

原始数据 #

import pandas as pd

df = pd.DataFrame({'ID': ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J'],
 'VAL': [9, 5, 8, 2, 6, 1, 1, 1, 2, 11]})

print(df)

输出：

ID  VAL
0  A    9
1  B    5
2  C    8
3  D    2
4  E    6
5  F    1
6  G    1
7  H    1
8  I    2
9  J   11

分组需求 #

创建GROUP列，需满足以下规则：

• 每组最多包含3行数据
• 组内VAL列总和不超过10
• 满足任一条件（行数达3或总和超10）则切换至下一组

预期输出：

ID  VAL  GROUP
0  A    9      1
1  B    5      2
2  C    8      3
3  D    2      3
4  E    6      4
5  F    1      4
6  G    1      4
7  H    1      5
8  I    2      5
9  J   11      6

现有尝试及问题 #

你尝试用全局累积求和和固定行号分组，但结果分组数超出预期：

identifier1 = df['VAL'].cumsum() // 10
identifier2 = [i // 3 for i in range(len(df))]

df['GROUP'] = df.groupby([identifier1, identifier2]).ngroup() + 1

print(df)

输出：

ID  VAL  GROUP
0  A    9      1
1  B    5      2
2  C    8      3
3  D    2      4
4  E    6      5
5  F    1      5
6  G    1      6
7  H    1      6
8  I    2      6
9  J   11      7

修正方案 #

问题核心在于你的方法用了全局规则（全局累积求和、固定行号整除），没有动态跟踪每组的实时状态（当前组的总和、行数）。需要逐行判断是否需要新建组，以下是两种可行方案：

方案1：循环逐行处理 #

current_group = 1
current_sum = 0
current_count = 0
groups = []

for val in df['VAL']:
    # 检查加入当前值后是否违反分组规则
    if current_sum + val > 10 or current_count >= 3:
        current_group += 1
        current_sum = val
        current_count = 1
    else:
        current_sum += val
        current_count += 1
    groups.append(current_group)

df['GROUP'] = groups
print(df)

输出与预期完全一致。

方案2：用itertools.accumulate实现（更简洁） #

通过accumulate传递当前组的状态（组号、总和、行数），动态更新分组：

from itertools import accumulate

def update_group(state, val):
    group_num, current_sum, current_count = state
    if current_sum + val > 10 or current_count >= 3:
        return (group_num + 1, val, 1)
    else:
        return (group_num, current_sum + val, current_count + 1)

# 初始状态：组号1、总和0、行数0
states = accumulate(df['VAL'], update_group, initial=(1, 0, 0))
# 跳过初始状态，提取每组的组号
df['GROUP'] = [s[0] for s in list(states)[1:]]
print(df)

同样能得到符合要求的分组结果。

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