优化时间序列DataFrame列下一次价格变动判断函数的方法

我有一个带价格列的大型DataFrame，价格会在一段时间内保持稳定，之后涨跌变动，再维持新价格一段时间。需要生成名为next movement的新列，标记下一次价格变动是上涨（1）、下跌（0）还是未知（-1）。

示例：
价格列：[1,1,1,2,2,2,4,4,4,3,3,3,4,4,4,2,1]
对应next movement列：[1,1,1,1,1,1,0,0,0,1,1,1,0,0,0,0,-1]

当前实现的函数功能可用但效率极低：

def make_next_movement_column(DataFrame, column):
    DataFrame["next movement"] = -1
    for i in range (DataFrame.shape[0]):
        for j in range(i + 1, DataFrame.shape[0]):
            if(DataFrame[column][j] > DataFrame[column][i]):
                DataFrame["next movement"][i:j] = 1
                break;
            if(DataFrame[column][j] < DataFrame[column][i]):
                DataFrame["next movement"][i:j] = 0
                break;
        i = j - 1
    return DataFrame

注：仅查看相邻列的diff方法方案不适用，我需要找到无论多远的下一次变动。

高效实现方案 #

利用pandas的向量化操作替代嵌套循环，可大幅提升处理大型数据集的效率，具体步骤如下：

• 定位价格变动点：找出价格发生变化的索引位置，同时补充最后一行索引以覆盖所有区间
• 计算变动方向：遍历每个价格稳定区间，确定该区间之后的下一次价格变动方向
• 映射到全量数据：将每个区间的方向值批量填充到对应行中

代码实现 #

import pandas as pd

def make_next_movement_column(df, price_col):
    # 复制原DataFrame，避免修改输入数据
    df = df.copy()
    # 初始化结果列为-1（未知）
    df["next movement"] = -1
    
    # 找出所有价格变动的临界点：当前行与前一行价格不同的索引
    change_points = df[df[price_col] != df[price_col].shift(1)].index.tolist()
    # 确保最后一行被包含，处理最后一个稳定区间
    if change_points[-1] != df.index[-1]:
        change_points.append(df.index[-1])
    
    # 遍历每个稳定区间，设置对应的下一次变动方向
    for i in range(len(change_points) - 1):
        start_idx = change_points[i]
        end_idx = change_points[i + 1]
        current_price = df.loc[start_idx, price_col]
        next_price = df.loc[end_idx, price_col]
        
        # 判断下一次变动方向并批量赋值
        if next_price > current_price:
            df.loc[start_idx:end_idx - 1, "next movement"] = 1
        elif next_price < current_price:
            df.loc[start_idx:end_idx - 1, "next movement"] = 0
    
    return df

代码说明 #

• 用shift()和布尔索引快速定位变动点，属于向量化操作，比逐行循环效率高几个数量级
• 仅遍历价格变动区间（数量远小于数据总行数），避免了嵌套循环的O(n²)时间复杂度
• 使用loc进行区间赋值，既避免链式索引的性能问题，也能防止SettingWithCopyWarning

测试验证 #

用示例数据测试，结果与预期一致：

test_df = pd.DataFrame({"price": [1,1,1,2,2,2,4,4,4,3,3,3,4,4,4,2,1]})
result_df = make_next_movement_column(test_df, "price")
print(result_df["next movement"].tolist())
# 输出：[1,1,1,1,1,1,0,0,0,1,1,1,0,0,0,0,-1]

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