如何按另一列指定大小将DataFrame分组拆分为数据块？

原始数据 #

import pandas as pd

data = {
    'A': ['car', 'car', 'car', 'car', 'truck', 'truck', 'truck'],
    'B': [2, 2, 2, 2, 1, 1, 1]
}

df = pd.DataFrame(data)

解决方案代码 #

def split_group(group):
    chunk_size = group['B'].iloc[0]
    group['C'] = (group.reset_index(drop=True).index // chunk_size).astype(int)
    return group

df = df.groupby('A', group_keys=False).apply(split_group)

运行结果 #

print(df)
# 输出：
#        A  B  C
# 0    car  2  0
# 1    car  2  0
# 2    car  2  1
# 3    car  2  1
# 4  truck  1  0
# 5  truck  1  1
# 6  truck  1  2

逻辑说明 #

• 按A列分组，group_keys=False确保分组后不会保留分组索引作为额外列
• 对每个分组执行自定义逻辑：
  1. 提取当前分组的拆分行数（即B列的值，题目中每组内B值统一，取第一个值即可）
  2. 重置分组内的行索引为从0开始，通过整数除法//按指定行数拆分，生成连续的组编号
  3. 将组编号存入新列C，返回处理后的分组

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