基于连续Treatment值对DataFrame资产进行分组的技术需求

示例数据 #

先还原你提供的DataFrame结构：

目标分组结果：

实现代码（Pandas） #

import pandas as pd

# 构造示例数据
data = {
    'Asset': ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K'],
    'Treatment': [1, 1, None, None, None, 2, None, None, None, None, 3]
}
df = pd.DataFrame(data)

# 1. 标记存在Treatment的行
df['has_treatment'] = df['Treatment'].notna()

# 2. 生成初始组号：每遇到有Treatment的行，组号递增
df['Group'] = df['has_treatment'].cumsum()

# 3. 向前填充组号，最多覆盖后续2个无Treatment的资产
df['Group'] = df['Group'].ffill(limit=2)

# 4. 把超量的无Treatment资产的组号设为空
df.loc[~df['has_treatment'] & df['Group'].isna(), 'Group'] = None

# 清理临时列，调整格式
df = df.drop('has_treatment', axis=1)
df['Group'] = df['Group'].astype('Int64')  # 保留整数型空值

print(df)

代码说明 #

• 标记有效Treatment行：用notna()识别非空Treatment记录，作为分组的起始标记
• 生成初始组号：cumsum()给每个连续的Treatment区块分配递增组号
• 限制填充范围：ffill(limit=2)确保只有组后前2个无Treatment资产继承组号
• 清空超量空值组号：对超过2个的无Treatment资产，将其组号置空，符合规则要求

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