如何用Python Pandas和Numpy基于多条件生成State变量？

已知如下Pandas DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'Value': [10,20,15,25,15,10,15],
    'Max(3)': [np.nan, np.nan, np.nan, 20,25,25,20],
    'Min(3)': [np.nan, np.nan, np.nan, 10,15,15,10]
})

需要计算State列，规则如下：

• a) 若Value > Max(3)，则State=1
• b) 若Value < Min(3)，则State=4
• c) 若Value处于[Min(3), Max(3)]区间，且前一个State为1或2，则State=2
• d) 若Value处于[Min(3), Max(3)]区间，且前一个State为4或3，则State=3
• 前3行因无有效Max(3)/Min(3)值，State设为nan

用户尝试用np.where()实现，但无法处理依赖前序State的逻辑，寻求解决方案。

核心原因：为什么np.where不适用？ #

np.where()是无状态的向量化操作，会一次性计算所有位置的结果，无法引用同一列中正在计算的前一个值。而你的需求属于有状态的序列计算——当前State的取值依赖上一行的结果，这种场景下必须用迭代或带状态的计算逻辑。

Excel里操作简单是因为单元格公式可以直接引用前一行单元格，本质是隐式的逐行迭代，只是Excel帮你封装了这个过程。

解法1：逐行迭代（直观易懂，适合小数据集） #

直接初始化状态变量，遍历每一行逐步计算：

# 初始化State列为nan
df['State'] = np.nan
prev_state = None

for idx in df.index:
    # 跳过前3行无有效窗口的记录
    if idx < 3:
        continue
    
    val = df.loc[idx, 'Value']
    max3 = df.loc[idx, 'Max(3)']
    min3 = df.loc[idx, 'Min(3)']
    
    if val > max3:
        curr_state = 1
    elif val < min3:
        curr_state = 4
    else:
        # 根据前一个状态判断当前值
        if prev_state in (1, 2):
            curr_state = 2
        elif prev_state in (3, 4):
            curr_state = 3
        else:
            curr_state = np.nan
    
    df.loc[idx, 'State'] = curr_state
    prev_state = curr_state

运行后得到的结果与期望一致：

Value  Max(3)  Min(3)  State
0     10     NaN     NaN    NaN
1     20     NaN     NaN    NaN
2     15     NaN     NaN    NaN
3     25    20.0    10.0    1.0
4     15    25.0    15.0    2.0
5     10    25.0    15.0    4.0
6     15    20.0    10.0    3.0

解法2：Numba加速（适合大数据集） #

如果数据集规模较大，逐行迭代效率偏低，可以用numba将循环编译为机器码，大幅提升速度：

首先安装依赖：

pip install numba

然后编写代码：

from numba import jit
import numpy as np

@jit(nopython=True)
def compute_state(values, max3, min3):
    n = len(values)
    state = np.full(n, np.nan)
    prev_state = np.nan
    
    for i in range(3, n):
        val = values[i]
        mx = max3[i]
        mn = min3[i]
        
        if val > mx:
            curr = 1
        elif val < mn:
            curr = 4
        else:
            if prev_state in (1, 2):
                curr = 2
            elif prev_state in (3, 4):
                curr = 3
            else:
                curr = np.nan
        state[i] = curr
        prev_state = curr
    return state

# 将DataFrame列转为numpy数组传入函数
df['State'] = compute_state(df['Value'].values, df['Max(3)'].values, df['Min(3)'].values)

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