如何用Pandas/PySpark创建span_flag_1与span_flag_2列？

我们可以使用Python的Pandas库来实现需求，核心思路是先将跨年的周数转换为连续数值，再针对每个flag列计算最近的上/下一次flag=1的周数，最后按规则生成span列。

步骤说明 #

1. 构造连续周数：将year和week转换为连续的累计周数，避免跨年导致的周数计算错误。
2. 定义通用计算函数：封装针对单个flag列的span计算逻辑，包括提取flag=1的周数、匹配最近的上/下一次周数、计算最大差值。
3. 应用函数到目标列：将函数分别应用到flag_1和flag_2，生成对应的span_flag_1和span_flag_2。

代码实现 #

import pandas as pd

# 构造样本数据
data = {
    'week': [26,27,28,2,3,4,5,6,7,8,9,10,11],
    'year': [2022,2022,2022,2023,2023,2023,2023,2023,2023,2023,2023,2023,2023],
    'flag_1': [0,1,0,0,1,0,1,0,0,0,0,0,0],
    'flag_2': [0,0,0,1,0,0,1,1,0,0,0,1,1]
}
df = pd.DataFrame(data)

# 生成连续周数：以2000年为基准计算累计周数，保证时间顺序连续性
df['continuous_week'] = (df['year'] - 2000) * 52 + df['week']

def calculate_span(df, flag_col):
    # 提取当前flag列所有值为1的连续周数
    flag_active_weeks = df[df[flag_col] == 1]['continuous_week'].tolist()
    
    # 匹配每行最近的上一次flag=1的周数
    df['last_active'] = df['continuous_week'].apply(
        lambda x: max([w for w in flag_active_weeks if w <= x], default=None)
    )
    # 匹配每行最近的下一次flag=1的周数
    df['next_active'] = df['continuous_week'].apply(
        lambda x: min([w for w in flag_active_weeks if w >= x], default=None)
    )
    
    # 初始化span列：flag=1时取值为1
    span_col_name = f'span_{flag_col}'
    df[span_col_name] = 1
    
    # 处理flag≠1的行：计算与上/下一次flag=1的周数差，取最大值
    non_flag_mask = df[flag_col] != 1
    df.loc[non_flag_mask, span_col_name] = df.loc[non_flag_mask].apply(
        lambda row: max(row['continuous_week'] - row['last_active'], row['next_active'] - row['continuous_week'])
        if row['last_active'] is not None and row['next_active'] is not None
        else (row['continuous_week'] - row['last_active'] if row['next_active'] is None else row['next_active'] - row['continuous_week'])
    )
    
    # 清理临时辅助列
    df.drop(['last_active', 'next_active'], axis=1, inplace=True)
    return df

# 分别计算两个flag对应的span列
df = calculate_span(df, 'flag_1')
df = calculate_span(df, 'flag_2')

# 输出结果（可移除continuous_week列展示最终数据）
print(df.drop('continuous_week', axis=1))

结果说明 #

执行代码后会得到包含span_flag_1和span_flag_2的完整数据集，例如：

• 2022年第28周（flag_1=0），最近上一次flag_1=1是第27周（差1），下一次是2023年第3周（差27），所以span_flag_1取最大值27。
• 2023年第7周（flag_2=0），最近上一次flag_2=1是第6周（差1），下一次是第10周（差3），所以span_flag_2取最大值3。

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