Python：为DataFrame基于条件新增exit_6months列的实现方案

原始数据 #

index    group      date    exit 
  0        A       Jan-22     0  
  1        A       Feb-22     0 
  2        A       Mar-22     0 
  3        A       Apr-22     0 
  5        A       May-22     0 
  6        A       Jun-22     0 
  7        A       Jul-22     0 
  8        A       Aug-22     0 
  9        A       Sep-22     1 
  6        B       Mar-21     0 
  7        B       Apr-21     0 
  8        B       May-21     0 
  9        B       Jun-21     0 
  10       B       Jul-21     0 
  11       B       Aug-21     0 
  12       B       Sep-21     0 
  13       B       Oct-21     1 
  14       C       Jan-23     0 
  15       C       Feb-23     0 
  16       C       Mar-23     1 
  17       D       Jul-23     0 
  18       D       Aug-23     1 

需求说明 #

新增exit_6months列，规则如下：

• 客户实际流失当月及流失前6个月的行标记为1
• 若客户流失前的记录不足6个月，则所有行均标记为1

预期结果 #

index    group      date    exit   exit_6months  
  0        A       Jan-22     0      0 
  1        A       Feb-22     0      0
  2        A       Mar-22     0      0
  3        A       Apr-22     0      1
  5        A       May-22     0      1
  6        A       Jun-22     0      1
  7        A       Jul-22     0      1
  8        A       Aug-22     0      1
  9        A       Sep-22     1      1
  6        B       Mar-21     0      0
  7        B       Apr-21     0      0
  8        B       May-21     0      1
  9        B       Jun-21     0      1
  10       B       Jul-21     0      1
  11       B       Aug-21     0      1
  12       B       Sep-21     0      1
  13       B       Oct-21     1      1
  14       C       Jan-23     0      1
  15       C       Feb-23     0      1
  16       C       Mar-23     1      1
  17       D       Jul-23     0      1
  18       D       Aug-23     1      1

实现方案（Python Pandas） #

步骤说明 #

1. 将日期列转换为标准datetime格式，并统一为月末日期，避免计算误差
2. 按客户分组，获取每个客户的流失日期
3. 计算每条记录日期与流失日期的月份间隔
4. 根据间隔月份判断是否标记为1

完整代码 #

import pandas as pd
from dateutil.relativedelta import relativedelta

# 构造原始数据（实际场景可替换为读取文件）
data = pd.DataFrame({
    'index': [0,1,2,3,5,6,7,8,9,6,7,8,9,10,11,12,13,14,15,16,17,18],
    'group': ['A']*9 + ['B']*8 + ['C']*3 + ['D']*2,
    'date': ['Jan-22','Feb-22','Mar-22','Apr-22','May-22','Jun-22','Jul-22','Aug-22','Sep-22',
             'Mar-21','Apr-21','May-21','Jun-21','Jul-21','Aug-21','Sep-21','Oct-21',
             'Jan-23','Feb-23','Mar-23','Jul-23','Aug-23'],
    'exit': [0]*8 + [1] + [0]*7 + [1] + [0]*2 + [1] + [0,1]
})

# 转换日期为月末datetime格式
data['date'] = pd.to_datetime(data['date'], format='%b-%y') + pd.offsets.MonthEnd(0)

# 为每个分组添加流失日期
def get_exit_date(group):
    exit_rows = group[group['exit'] == 1]
    return exit_rows['date'].iloc[0] if not exit_rows.empty else None

data['exit_date'] = data.groupby('group')['date'].transform(get_exit_date)

# 计算当前日期到流失日期的月份差
data['months_to_exit'] = data.apply(
    lambda row: relativedelta(row['exit_date'], row['date']).months,
    axis=1
)

# 生成目标列
data['exit_6months'] = data['months_to_exit'].apply(lambda x: 1 if x <= 6 else 0)

# 整理结果格式，还原日期显示
result = data[['index', 'group', 'date', 'exit', 'exit_6months']]
result['date'] = result['date'].dt.strftime('%b-%y')

# 打印结果
print(result.to_string(index=False))

关键细节 #

• 使用pd.offsets.MonthEnd(0)确保日期统一为月末，避免因不同月份天数导致的计算偏差
• 通过groupby.transform为每个客户的所有记录批量添加流失日期，无需循环处理
• relativedelta准确计算月份间隔，不受跨年度、不同月份天数影响
• 不足6个月的客户自然满足months_to_exit <=6，自动标记为1，无需额外判断

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