如何为降雨数据实现延长最长干旱期1天/2天的两种场景

• 基于存储降雨数据的DataFrame，实现函数支持两种场景：
  • 场景A：将最长干旱期（降雨量≤1mm转为0后，连续无降雨的天数）延长1天
  • 场景B：将该最长干旱期延长2天
• 需返回处理后的DataFrame、干旱期起止时间、时长，支持单独获取单场景结果或同时获取两个场景结果
• 当前已完成场景A代码，需整合场景B逻辑并优化参数设计

核心思路 #

1. 分离通用逻辑与修改逻辑：先找到最长干旱期的起止索引（不修改数据），再根据需要的延长天数修改DataFrame，避免修改过程干扰干旱期判断
2. 支持灵活参数：允许传入单个延长天数（如1或2）或天数列表（如[1,2]），返回对应格式的结果
3. 保护原始数据：操作前复制DataFrame，避免修改输入的原始数据

完整代码实现 #

def increase_dry_period(df, extend_days):
    # 复制原始DataFrame，避免修改原数据
    df_processed = df.copy()
    
    # 第一步：将降雨量≤1mm的记录转为0，统一干旱判断标准
    idx_zero = df_processed['PRECIP'] <= 1
    df_processed.loc[idx_zero, 'PRECIP'] = 0
    
    season_length = df_processed.shape[0]
    max_counter = 0
    counter = 0
    start_row = 0
    end_row = 0
    current_start = 0
    
    # 第二步：遍历数据，定位最长干旱期的起止索引
    for k, row in df_processed.iterrows():
        P = row['PRECIP']
        if P == 0:
            # 记录当前干旱期的起始索引
            if counter == 0:
                current_start = k
            counter += 1
            # 更新最长干旱期的起止信息
            if counter > max_counter:
                max_counter = counter
                start_row = current_start
                end_row = k
        else:
            # 遇到降雨，重置当前干旱期计数器
            counter = 0
    
    # 封装延长干旱期的逻辑，避免重复代码
    def apply_extension(days):
        df_copy = df_processed.copy()
        if end_row < season_length - 1:
            # 干旱期未到季末，向后延长：修改干旱期结束后的days天为0
            for i in range(1, days + 1):
                if end_row + i < season_length:
                    df_copy.loc[end_row + i, 'PRECIP'] = 0.0
            new_end = min(end_row + days, season_length - 1)
            final_start = start_row
            final_end = new_end
        else:
            # 干旱期已到季末，向前延长：修改干旱期开始前的days天为0
            for i in range(1, days + 1):
                if start_row - i >= 0:
                    df_copy.loc[start_row - i, 'PRECIP'] = 0.0
            new_start = max(start_row - days, 0)
            final_start = new_start
            final_end = end_row
        
        # 计算延长后的干旱期时长（与原代码逻辑保持一致）
        new_length = final_end - final_start
        
        return {
            'df': df_copy,
            'start': df_copy.loc[final_start, 'TIMESTAMP'],
            'end': df_copy.loc[final_end, 'TIMESTAMP'],
            'length': new_length
        }
    
    # 根据传入的参数类型返回对应结果
    if isinstance(extend_days, int):
        return apply_extension(extend_days)
    elif isinstance(extend_days, list):
        results = {}
        for days in extend_days:
            results[f'scenario_{days}'] = apply_extension(days)
        return results
    else:
        raise ValueError("extend_days参数必须是整数或整数列表")

使用示例 #

# 场景A：获取延长1天的结果
result_a = increase_dry_period(original_df, 1)
print("场景A干旱期起止：", result_a['start'], "-", result_a['end'])
print("场景A干旱期时长：", result_a['length'])
processed_df_a = result_a['df']

# 场景B：获取延长2天的结果
result_b = increase_dry_period(original_df, 2)

# 同时获取两个场景的结果
all_results = increase_dry_period(original_df, [1, 2])
scenario_a = all_results['scenario_1']
scenario_b = all_results['scenario_2']

代码优化说明 #

• 修复原代码逻辑问题：原代码在遍历过程中直接修改DataFrame，会干扰后续干旱期的判断，改为先定位最长干旱期再修改，结果更准确
• 保护原始数据：通过df.copy()避免修改输入的原始DataFrame
• 逻辑复用：用内部函数封装延长逻辑，减少重复代码，便于后续扩展更多延长天数场景
• 灵活参数：支持单个或多个延长天数，满足不同场景需求

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