使用Pandas GroupBy从单列生成多列：按0分割折扣列

需求说明 #

给定包含date和discount列的DataFrame，需将discount列中的非零值以单个或连续的0为分隔符分组，每组对应生成一列split_n，非对应组的位置填充NaN，零值对应的所有split列均为NaN。

示例输入 #

import pandas as pd
import numpy as np

# 构造示例输入
df = pd.DataFrame({
    'date': pd.date_range(start='2023-01-01', periods=8),
    'discount': [1, 2, 0, 3, 0, 0, 4, 5]
})

目标输出 #

正确实现代码 #

# 1. 为非零值分配组号：遇到"0后接非零"的情况时组号递增
df['group'] = (df['discount'] != 0) & (df['discount'].shift(1) == 0)
df['group'] = df['group'].cumsum()
# 零值的组号设为NaN，避免干扰后续转列
df.loc[df['discount'] == 0, 'group'] = np.nan

# 2. 将组号转成列，生成split系列列
split_cols = df.pivot(columns='group', values='discount')
split_cols.columns = [f'split{int(col)}' for col in split_cols.columns]

# 3. 合并回原DataFrame并清理临时列
result_df = pd.concat([df.drop('group', axis=1), split_cols], axis=1)

关键逻辑说明 #

• 分组标识通过判断当前值非零且前一个值为零来触发组号递增，确保连续非零值归为同一组，连续零仅作为分隔符。
• 使用pivot方法将分组转成列，自动为非对应组的位置填充NaN，完美匹配需求。

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