如何加速Python中处理DataFrame的For循环？业务场景优化求助

核心问题分析 #

你的原代码用了两层嵌套循环（外层for+内层while），时间复杂度是O(n²)，百万级数据下会出现指数级的耗时增长。要实现秒级/分钟级处理，必须把复杂度降到O(n)。

优化思路：从后往前的动态规划遍历 #

我们可以反向遍历价格序列，维护两个数组记录每个位置往后第一个触发上涨突破（≥当前价+2）和下跌突破（≤当前价-2）的索引，最后通过比较这两个索引的先后，确定每个位置的updown标记。这种方法只需要遍历序列两次（一次反向初始化，一次正向判断），时间复杂度为O(n)。

优化后的代码 #

import pandas as pd
import numpy as np

# 示例数据
df = pd.DataFrame.from_dict({'price': {0: 98, 1: 99, 2: 101, 3: 99, 4: 97, 5: 100, 6: 100, 7: 98}})
prices = df['price'].to_numpy()
n = len(prices)

# 初始化两个数组，用n（超出序列长度）表示无后续突破
last_up = np.full(n, n)
last_down = np.full(n, n)

# 从倒数第二个元素开始反向遍历
for i in range(n-2, -1, -1):
    current_price = prices[i]
    next_price = prices[i+1]
    
    # 更新第一个上涨突破的索引
    if next_price >= current_price + 2:
        last_up[i] = i + 1
    elif next_price <= current_price - 2:
        # 下一个元素已触发下跌，继承下一个元素的上涨突破索引
        last_up[i] = last_up[i+1]
    else:
        # 下一个元素未突破，直接继承其上涨突破索引
        last_up[i] = last_up[i+1]
    
    # 更新第一个下跌突破的索引
    if next_price <= current_price - 2:
        last_down[i] = i + 1
    elif next_price >= current_price + 2:
        # 下一个元素已触发上涨，继承下一个元素的下跌突破索引
        last_down[i] = last_down[i+1]
    else:
        # 下一个元素未突破，直接继承其下跌突破索引
        last_down[i] = last_down[i+1]

# 生成updown列
df['updown'] = np.where(last_up < last_down, 'Up', 
                        np.where(last_down < last_up, 'Down', 0))

关键优化点说明 #

1. 用NumPy数组替代Pandas逐行访问：NumPy的数组访问速度远快于df.price.iloc[i]，避免了Pandas的行索引开销。
2. O(n)时间复杂度：反向遍历只需要一次线性扫描，彻底摆脱了原代码的嵌套循环。
3. 向量化判断：最后用np.where生成结果列，比Python循环更快。

学习途径建议 #

• 先理解时间复杂度的概念：明确O(n)、O(n²)等复杂度在大数据量下的差异，这是优化代码的核心思路。
• 熟练掌握Pandas/NumPy的向量化操作：尽量避免逐行循环，优先使用内置的向量化方法（如np.where、pd.Series.shift等）。
• 学习动态规划思想：这类“找后续第一个满足条件”的问题，反向遍历+状态继承是常用的高效解法。
• 可选进阶：如果遇到更复杂的循环场景，可以学习Numba工具，它能把Python循环编译成机器码，进一步提升速度。

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