Python DataFrame嵌套循环向量化优化：降低时间复杂度

需求回顾 #

• 基于两种时间约束场景，为Diff_*列生成新列：
  1. 每周7天，时间范围7:00-23:00，对应固定小时数5824
  2. 每周工作日（周一至周五），时间范围7:00-23:00，对应固定小时数4160
• 新列规则：对每个Diff_*列，计算列总和/对应小时数得到最大值，生成0到该最大值、步长为10的偏移量列表；在符合时间约束的行中，生成原列值 - 偏移量的新列

原代码的低效点 #

1. 循环求和：用for index in df_current.index累加求和，效率远低于pandas内置的sum()方法
2. 逐行判断赋值：嵌套循环遍历索引行进行条件判断和赋值，完全没有利用pandas的向量化特性
3. 时间匹配低效：用any(x in df_current['Time'][idx] for x in allowed_time)做字符串匹配，不如将Time列转为时间类型后直接判断小时范围
4. 重复逻辑冗余：两种场景代码高度重复，没有封装成可复用函数

优化后的代码实现 #

首先做数据预处理，把Time和Date列转为合适的类型，提前生成两种场景的布尔掩码：

import pandas as pd
import numpy as np

# 预处理：转换时间和日期列类型
df_current['Time'] = pd.to_datetime(df_current['Time']).dt.time
df_current['Date'] = pd.to_datetime(df_current['Date'])

# 生成时间范围掩码（7:00-23:00）
time_mask = df_current['Time'].between(pd.to_datetime('7:00:00').time(), pd.to_datetime('23:00:00').time())
# 生成工作日掩码（周一至周五，weekday()返回0-4）
weekday_mask = df_current['Date'].dt.weekday <= 4

# 封装通用处理函数
def generate_offset_columns(df, diff_cols, hours_constant, filter_mask):
    for col in diff_cols:
        # 向量化求和替代循环累加
        col_sum = df[col].sum()
        max_fixed = col_sum / hours_constant
        # 生成偏移量列表
        offsets = np.arange(0, round(max_fixed), 10)
        # 批量生成新列
        for offset in offsets:
            new_col_name = f"{col}_{offset}"
            # 初始化新列为原列值
            df[new_col_name] = df[col]
            # 仅在符合条件的行应用偏移
            df.loc[filter_mask, new_col_name] -= offset
    return df

# 场景1：7天+7:00-23:00
diff_cols_7d = [f'Diff_{y}' for y in reqd_list]
df_current = generate_offset_columns(df_current, diff_cols_7d, 5824, time_mask)

# 场景2：工作日+7:00-23:00
diff_cols_5d = df_current.iloc[:, 29:].columns.tolist()
# 组合两个掩码：工作日+时间范围
weekday_time_mask = weekday_mask & time_mask
df_current = generate_offset_columns(df_current, diff_cols_5d, 4160, weekday_time_mask)

优化关键点说明 #

1. 向量化运算：用df[col].sum()替代循环求和，用df.loc[mask, col]替代逐行赋值，这是提升效率的核心
2. 提前生成掩码：把时间和日期的条件判断提前转为布尔掩码，避免在循环中重复计算
3. 函数复用：把两种场景的共同逻辑封装成函数，减少代码冗余，也便于维护
4. 时间类型优化：将Time列转为time类型后，用between()做范围判断，比字符串匹配快得多

额外性能建议 #

• 如果Diff_*列数量极多，可以考虑用pd.concat批量生成新列，进一步减少循环层级
• 若DataFrame规模超大，可以尝试用Dask进行并行处理，但对于常规规模的DataFrame，上述优化已足够解决耗时问题

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