如何优化OHLC数据框判断后续价格先涨跌指定百分比的代码

原有代码问题 #

你现有的逐行apply方案时间复杂度为O(n²)，每次遍历都要做切片、布尔索引操作，纯Python循环的执行效率极低，数据量超过1万行时卡顿会非常明显，且原代码存在函数名笔误（定义的函数名为check，调用时写的是find_threshold）。

最优实现方案（Numba JIT加速） #

这是目前兼顾速度和内存开销的最佳方案，通过Numba将纯Python循环编译为机器码，执行速度比原方案提升100~1000倍，支持百万级OHLC数据正常运行：

依赖安装 #

如果没有装Numba可以先执行：pip install numba

代码实现 #

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

@jit(nopython=True)  # 开启nopython模式编译为机器码
def find_first_threshold(close_arr, high_arr, low_arr, threshold=0.01):
    n = len(close_arr)
    res = np.zeros(n, dtype=np.int8)  # 用int8节省内存
    for i in range(n):
        current_close = close_arr[i]
        up_thr = current_close * (1 + threshold)
        down_thr = current_close * (1 - threshold)
        first_up_idx = n  # 初始值设为数组长度，表示未找到
        first_down_idx = n
        # 从当前行的下一行开始遍历
        for j in range(i + 1, n):
            # 优先找最早触发的阈值
            if high_arr[j] >= up_thr and first_up_idx == n:
                first_up_idx = j
            if low_arr[j] <= down_thr and first_down_idx == n:
                first_down_idx = j
            # 两个阈值都找到就提前退出循环，减少无效计算
            if first_up_idx != n and first_down_idx != n:
                break
        # 赋值结果
        if first_up_idx < first_down_idx:
            res[i] = 1
        elif first_down_idx < first_up_idx:
            res[i] = -1
        else:
            res[i] = 0
    return res

# 调用函数直接生成结果列
ohlc['check'] = find_first_threshold(
    ohlc['close'].values,
    ohlc['high'].values,
    ohlc['low'].values
)

替代方案（无Numba纯Pandas向量化） #

如果无法安装Numba，且数据量小于1万行，可以用广播向量化实现，避免逐行apply，但数据量过大时会出现内存溢出：

import numpy as np
import pandas as pd

close = ohlc['close'].values
high = ohlc['high'].values
low = ohlc['low'].values
n = len(close)

# 生成上穿、下穿阈值的布尔矩阵，上三角矩阵仅保留当前行之后的判断
up_mask = np.triu(high > (close * 1.01).reshape(-1, 1), k=1)
down_mask = np.triu(low < (close * 0.99).reshape(-1, 1), k=1)

# 找每行第一个触发阈值的位置
first_up = np.argmax(up_mask, axis=1)
first_down = np.argmax(down_mask, axis=1)

# 处理未触发的情况（argmax默认返回0，需要替换为超出数组长度的值）
first_up[~up_mask.any(axis=1)] = n
first_down[~down_mask.any(axis=1)] = n

# 生成结果列
ohlc['check'] = np.where(
    first_up < first_down, 1,
    np.where(first_down < first_up, -1, 0)
)

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