Pandas中多分支IF处理多列生成新列报错求助

我有一个名为mrp的DataFrame，想要根据rem_dys列的值，通过多分支IF逻辑计算对应p系列列的累加和，生成daily_consumption新列。逻辑规则为：rem_dys列值<=30时累加p01+p02，<=60时累加p01到p03，以此类推，步长30天。编写的代码如下：

def planning_period(x):
    if x <= 30:
        return mrp['p01'] + mrp['p02']
    elif x <= 60:
        return mrp['p01'] + mrp['p02'] + mrp['p03']
    elif x <= 90:
        return mrp['p01'] + mrp['p02'] + mrp['p03'] + mrp['p04']
    elif x <= 120:
        return mrp['p01'] + mrp['p02'] + mrp['p03'] + mrp['p04'] + mrp['p05']
    elif x <= 150:
        return mrp['p01'] + mrp['p02'] + mrp['p03'] + mrp['p04'] + mrp['p05'] + mrp['p06']
    elif x <= 180:
        return mrp['p01'] + mrp['p02'] + mrp['p03'] + mrp['p04'] + mrp['p05'] + mrp['p06'] + mrp['p07']
    elif x <= 210:
        return mrp['p01'] + mrp['p02'] + mrp['p03'] + mrp['p04'] + mrp['p05'] + mrp['p06'] + mrp['p07'] + mrp['p08']
    elif x <= 240:
        return mrp['p01'] + mrp['p02'] + mrp['p03'] + mrp['p04'] + mrp['p05'] + mrp['p06'] + mrp['p07'] + mrp['p08'] + mrp['p09']
    elif x <= 270:
        return mrp['p01'] + mrp['p02'] + mrp['p03'] + mrp['p04'] + mrp['p05'] + mrp['p06'] + mrp['p07'] + mrp['p08'] + mrp['p09'] + mrp['p10']
    else:
        return mrp['p01'] + mrp['p02'] + mrp['p03'] + mrp['p04'] + mrp['p05'] + mrp['p06'] + mrp['p07'] + mrp['p08'] + mrp['p09'] + mrp['p10'] + mrp['p11']

mrp['daily_consumption'] = mrp['rem_dys'].apply(planning_period)

运行代码时出现报错，报错信息显示期望1D数组，却得到形状为(1688,1688)的数组，同时伴随KeyError。

• 维度不匹配：apply调用planning_period时，每次传入的是rem_dys的单个值，但函数返回的是整个DataFrame的列相加结果（一个Series），而非当前行对应的单个累加值。当apply尝试将每个返回的Series赋值给新列的单个元素时，就会触发维度错误。
• KeyError：大概率是DataFrame中不存在部分p系列列，或者列名拼写错误。

方法1：修正函数，处理单行数据 #

修改函数使其接收整行数据，针对当前行计算对应p列的累加和：

def planning_period(row):
    x = row['rem_dys']
    cols = []
    if x <= 30:
        cols = ['p01', 'p02']
    elif x <= 60:
        cols = ['p01', 'p02', 'p03']
    elif x <= 90:
        cols = ['p01', 'p02', 'p03', 'p04']
    elif x <= 120:
        cols = ['p01', 'p02', 'p03', 'p04', 'p05']
    elif x <= 150:
        cols = ['p01', 'p02', 'p03', 'p04', 'p05', 'p06']
    elif x <= 180:
        cols = ['p01', 'p02', 'p03', 'p04', 'p05', 'p06', 'p07']
    elif x <= 210:
        cols = ['p01', 'p02', 'p03', 'p04', 'p05', 'p06', 'p07', 'p08']
    elif x <= 240:
        cols = ['p01', 'p02', 'p03', 'p04', 'p05', 'p06', 'p07', 'p08', 'p09']
    elif x <= 270:
        cols = ['p01', 'p02', 'p03', 'p04', 'p05', 'p06', 'p07', 'p08', 'p09', 'p10']
    else:
        cols = ['p01', 'p02', 'p03', 'p04', 'p05', 'p06', 'p07', 'p08', 'p09', 'p10', 'p11']
    return row[cols].sum()

mrp['daily_consumption'] = mrp.apply(planning_period, axis=1)

• 关键改动：apply时指定axis=1，让函数接收每一行数据；函数内根据当前行的rem_dys值选择对应列，再计算该行这些列的和，返回单个数值，避免维度不匹配。

方法2：向量化处理（更高效） #

针对大数据集，向量化操作比逐行apply性能更优：

# 获取所有p系列列并按顺序排序
p_cols = sorted([col for col in mrp.columns if col.startswith('p')])
# 计算每行的累加和矩阵（从p01到p11逐步累加）
cumulative_sums = mrp[p_cols].cumsum(axis=1)

# 定义rem_dys区间与对应累加列索引的映射
def get_cumulative_index(x):
    if x <= 30:
        return 1  # 对应前2列（p01-p02）的累加和，索引从0开始
    elif x <= 60:
        return 2
    elif x <= 90:
        return 3
    elif x <= 120:
        return 4
    elif x <= 150:
        return 5
    elif x <= 180:
        return 6
    elif x <= 210:
        return 7
    elif x <= 240:
        return 8
    elif x <= 270:
        return 9
    else:
        return 10  # 对应前11列（p01-p11）的累加和

# 生成每行对应的累加列索引
indices = mrp['rem_dys'].apply(get_cumulative_index)
# 根据索引提取对应累加和
mrp['daily_consumption'] = cumulative_sums.lookup(cumulative_sums.index, cumulative_sums.columns[indices])

• 优势：避免逐行循环，处理大数据集时速度更快；逻辑更简洁，减少重复代码。

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