如何优雅实现基于前一行值更新DataFrame的position字段？

需求说明 #

需要根据当前行的factor值和前一行的position值，计算当前行的position值。原循环代码不仅繁琐，还无法正确更新数值，寻求更优实现。

初始数据示例 #

factor  position
time                                   
2022-05-13 06:00:00  0.489471         0
2022-05-13 07:00:00  0.711030         0
2022-05-13 08:00:00  0.566865         0
2022-05-13 09:00:00  0.489471         0
2022-05-13 10:00:00  0.288419         0

现有问题代码 #

import pandas as pd

df = pd.DataFrame({'time': ['2022-05-13 06:00:00', '2022-05-13 07:00:00', '2022-05-13 08:00:00','2022-05-13 09:00:00', '2022-05-13 10:00:00'],
                   'factor': [0.489471, 0.711030, 0.566865, 0.489471, 0.288419],
                   'position': [0, 0, 0, 0, 0]})
df['time'] = pd.to_datetime(df['time'])
df.set_index('time', inplace=True)

threshold_2 = 0.7
threshold_1 = 0.35

for i in range(0, len(df)):
    # no position
    if i == 0 or df.iloc[i-1, :]['position'] == 0:
        if df.iloc[i, :]['factor'] > threshold_2:
            df.iloc[i, :]['position'] = 1
        else:
            df.iloc[i, :]['position'] = 0

    #has position
    elif df.iloc[i-1, :]['position'] != 0:
        if df.iloc[i, :]['factor'] > threshold_1:
            df.iloc[i, :]['position'] = 1
        else:
            df.iloc[i, :]['position'] = 0

原代码问题分析 #

• 链式索引（df.iloc[i, :]['position']）会触发SettingWithCopyWarning，无法保证赋值操作作用于原DataFrame，导致更新失败。
• 循环遍历在大数据量场景下效率低下。

解决方案 #

方案1：修正循环赋值逻辑（小数据量适用） #

直接通过列索引定位赋值，避免链式索引，确保更新作用于原DataFrame：

import pandas as pd

df = pd.DataFrame({'time': ['2022-05-13 06:00:00', '2022-05-13 07:00:00', '2022-05-13 08:00:00','2022-05-13 09:00:00', '2022-05-13 10:00:00'],
                   'factor': [0.489471, 0.711030, 0.566865, 0.489471, 0.288419],
                   'position': [0, 0, 0, 0, 0]})
df['time'] = pd.to_datetime(df['time'])
df.set_index('time', inplace=True)

threshold_2 = 0.7
threshold_1 = 0.35

# 提前获取列索引，避免重复计算
pos_col = df.columns.get_loc('position')
factor_col = df.columns.get_loc('factor')

for i in range(len(df)):
    if i == 0 or df.iloc[i-1, pos_col] == 0:
        df.iloc[i, pos_col] = 1 if df.iloc[i, factor_col] > threshold_2 else 0
    else:
        df.iloc[i, pos_col] = 1 if df.iloc[i, factor_col] > threshold_1 else 0

方案2：用Numba加速循环（大数据量适用） #

对于百万级以上的数据集，用Numba编译循环可以大幅提升效率：

import pandas as pd
from numba import njit
import numpy as np

df = pd.DataFrame({'time': ['2022-05-13 06:00:00', '2022-05-13 07:00:00', '2022-05-13 08:00:00','2022-05-13 09:00:00', '2022-05-13 10:00:00'],
                   'factor': [0.489471, 0.711030, 0.566865, 0.489471, 0.288419],
                   'position': [0, 0, 0, 0, 0]})
df['time'] = pd.to_datetime(df['time'])
df.set_index('time', inplace=True)

threshold_2 = 0.7
threshold_1 = 0.35

# 用Numba编译状态依赖的计算逻辑
@njit
def calculate_position(factor_arr, initial_pos, thresh2, thresh1):
    n = len(factor_arr)
    pos_arr = np.zeros(n, dtype=np.int64)
    pos_arr[0] = 1 if factor_arr[0] > thresh2 else initial_pos
    for i in range(1, n):
        if pos_arr[i-1] == 0:
            pos_arr[i] = 1 if factor_arr[i] > thresh2 else 0
        else:
            pos_arr[i] = 1 if factor_arr[i] > thresh1 else 0
    return pos_arr

# 转换为numpy数组传入函数，提升效率
df['position'] = calculate_position(df['factor'].values, df['position'].iloc[0], threshold_2, threshold_1)

方案3：用pandas的apply结合状态变量（代码简洁） #

利用apply方法，通过外部变量保存前一行的状态：

import pandas as pd

df = pd.DataFrame({'time': ['2022-05-13 06:00:00', '2022-05-13 07:00:00', '2022-05-13 08:00:00','2022-05-13 09:00:00', '2022-05-13 10:00:00'],
                   'factor': [0.489471, 0.711030, 0.566865, 0.489471, 0.288419],
                   'position': [0, 0, 0, 0, 0]})
df['time'] = pd.to_datetime(df['time'])
df.set_index('time', inplace=True)

threshold_2 = 0.7
threshold_1 = 0.35

# 初始化前一行状态
prev_pos = df['position'].iloc[0]

def update_pos(row):
    global prev_pos
    if prev_pos == 0:
        current_pos = 1 if row['factor'] > threshold_2 else 0
    else:
        current_pos = 1 if row['factor'] > threshold_1 else 0
    prev_pos = current_pos
    return current_pos

df['position'] = df.apply(update_pos, axis=1)

验证结果 #

运行上述代码后，position列的最终结果为：

time
2022-05-13 06:00:00    0
2022-05-13 07:00:00    1
2022-05-13 08:00:00    1
2022-05-13 09:00:00    1
2022-05-13 10:00:00    0
Name: position, dtype: int64

内容的提问来源于stack exchange，提问作者Will-NotGiveUp