基于第三列切换变量，在Pandas中生成计算列

问题需求 #

基于DataFrame的load列计算新列calculated load，需根据position列的值动态切换计算所用变量：

• 初始状态使用s_rising变量，公式为：load + x * s_rising
• 当position列出现值1时，后续行切换为s_falling变量，公式为：load + x * s_falling
• 当position列出现值-1时，后续行换回s_rising变量计算

解决方案代码 #

import pandas as pd
import numpy as np

s_falling = -4
s_rising = 2
x = 2

df = pd.DataFrame({"load": [1,2,4,6,2,4,7,4,8,3,4,7,3,3,6,4,7,4,3,2],
                   "position": [0,0.2,0.5,0.8,0.7,1,0.7,0.6,0.7,0.8,0.4,0.2,0,-0.5,-0.8,-1,-0.8,-0.9,-0.7,-0.6]})

# 生成切换信号：以上一行的position值判断是否切换状态
signal = np.where(df['position'].shift() == 1, 1, 
                  np.where(df['position'].shift() == -1, -1, 0))

# 计算当前状态：累积和确保状态在0（rising）和1（falling）间切换
state = signal.cumsum().clip(lower=0)

# 根据状态选择对应的计算变量
selected_s = np.where(state > 0, s_falling, s_rising)

# 计算目标列
df['calculated load'] = df['load'] + x * selected_s

print(df['calculated load'])

代码逻辑说明 #

1. 切换信号生成：使用shift()获取上一行的position值，标记当前行是否需要切换计算变量——上一行是1则标记为切换到falling状态，上一行是-1则标记为切换回rising状态。
2. 状态跟踪：通过累积和计算当前所处的状态，clip(lower=0)保证状态不会出现负数，仅在0（rising）和1（falling）两种状态间切换。
3. 变量选择与计算：根据当前状态选择对应的s变量，代入公式生成最终的calculated load列。

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