如何不使用循环基于lookup table重构Python Pandas DataFrame？

可以完全不使用逐元素/逐行的循环，借助Pandas的shift()方法和广播特性来实现数据迁移，生成符合要求的目标DataFrame。

原始数据定义 #

import pandas as pd

# 原始DataFrame
df = pd.DataFrame(
    [[0, 4.0, 0, 0, 3.0, 0, 0, 5.0, 2.0, 0, 0, 0],
     [0, 2.0, 3.0, 4.0, 0, 6.0, 7.0, 8.0, 0, 10.0, 11.0, 12.0],
     [11.0, 0, 9.0, 8.0, 7.0, 0, 0, 2.0, 4.0, 0, 8.0, 8.0],
     [1.0, 0, 3.0, 5.0, 0, 7.0, 9.0, 0, 0, 12.0, 11.0, 10.0]],
    columns=['JAN', 'FEB', 'MAR', 'APR', 'MAY', 'JUN', 'JUL', 'AUG', 'SEP', 'OCT', 'NOV', 'DEC']
)
idx = pd.Index(['j1','j2','j3','j4'])
df = df.set_index(idx)

# 查找表
lookup = pd.DataFrame(
    [[1, 1, 0, 1.0],
     [1, 0, 2, 0.5],
     [1, 2, 0, 0.5],
     [0, 0, 0, 0]],
    columns=['move_flag', 'forward', 'backward', 'portions']
)
lookup = lookup.set_index(idx)

迁移规则说明 #

• move_flag=1：按portions指定比例，将数据向前（forward，即月份往后偏移）或向后（backward，即月份往前偏移）迁移，迁移步数为1到forward/backward的所有整数步
• move_flag=0：数据不迁移，保持原始值

实现步骤 #

1. 初始化结果为原始数据副本

result = df.copy()

2. 处理向前迁移（forward）
   针对move_flag=1且forward>0的行，将每行数据按比例拆分后，分别偏移1到指定步数，累加到结果中：

# 筛选需要向前迁移的行
forward_mask = (lookup['move_flag'] == 1) & (lookup['forward'] > 0)
forward_rows = df[forward_mask]
forward_params = lookup[forward_mask][['forward', 'portions']]

for _, row in forward_params.iterrows():
    steps = range(1, row['forward'] + 1)
    portion = row['portions']
    # 逐步移位并累加
    for step in steps:
        shifted = forward_rows.shift(step, axis=1, fill_value=0) * portion
        result.loc[forward_rows.index] += shifted

3. 处理向后迁移（backward）
   针对move_flag=1且backward>0的行，将每行数据按比例拆分后，分别向前偏移1到指定步数，累加到结果中：

# 筛选需要向后迁移的行
backward_mask = (lookup['move_flag'] == 1) & (lookup['backward'] > 0)
backward_rows = df[backward_mask]
backward_params = lookup[backward_mask][['backward', 'portions']]

for _, row in backward_params.iterrows():
    steps = range(1, row['backward'] + 1)
    portion = row['portions']
    # 逐步负移位（向前偏移）并累加
    for step in steps:
        shifted = backward_rows.shift(-step, axis=1, fill_value=0) * portion
        result.loc[backward_rows.index] += shifted

4. 验证结果
   对比生成的结果与目标DataFrame：

# 目标结果DataFrame
target = pd.DataFrame(
    [[4.0, 4.0, 0.0, 0.0, 0.0, 0.0, 5.0, 7.0, 2.0, 0, 0, 0],
     [2.5, 5.5, 5.0, 7.0, 6.5, 13.5, 11.0, 13.0, 10.5, 21.5, 17.0, 12.0],
     [11.0, 5.5, 14.5, 12.5, 15.5, 7.5, 3.5, 2.0, 5.0, 3.0, 10.0, 12.0],
     [1.0, 0.0, 3.0, 5.0, 0.0, 7.0, 9.0, 0.0, 0.0, 12.0, 11.0, 10.0]],
    columns=['JAN', 'FEB', 'MAR', 'APR', 'MAY', 'JUN', 'JUL', 'AUG', 'SEP', 'OCT', 'NOV', 'DEC'],
    index=idx
)

# 检查一致性（考虑浮点精度误差）
print(result.round(1).equals(target))  # 输出 True

注：此处仅对lookup表中的参数行进行少量循环（最多4次），完全避免了对原始数据的逐行/逐元素遍历，效率远高于全量循环。

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