如何简化Python中DataFrame的GapUp/GapDown判断代码？

需求说明 #

对比DataFrame中当前行的Open值与前一行的Close值，当差值超过指定阈值（示例为3）时，在Gap列标记GapUp（向上跳空）或GapDown（向下跳空），第一行无前置数据，标记为NaN。

原循环方案的问题 #

原代码采用逐行循环的方式，存在以下不足：

• 效率低下，数据量较大时运行速度慢
• 需要手动处理索引范围，容易出现边界异常
• 代码冗余，可读性较差

矢量化优化方案 #

利用Pandas的矢量化操作和Numpy的条件判断，可实现简洁高效的批量处理，无需手动循环：

import pandas as pd
import numpy as np

# 若Open/Close为带逗号的字符串格式，先转换为数值类型
df['Open'] = df['Open'].str.replace(',', '.').astype(float)
df['Close'] = df['Close'].str.replace(',', '.').astype(float)

# 计算当前Open与前一行Close的差值
N = 3
diff = df['Open'] - df['Close'].shift(1)

# 批量标记Gap类型
df['Gap'] = np.where(diff > N, 'GapUp',
                     np.where(diff < -N, 'GapDown', np.nan))

方案优势 #

1. 性能高效：矢量化操作基于底层优化，比逐行循环快数十倍，尤其适合大数据集
2. 代码简洁：无需处理索引循环，一行代码完成条件判断与赋值
3. 自动处理边界：shift(1)会自动让第一行的差值为NaN，对应Gap列自动设为NaN，无需额外初始化处理

验证结果 #

处理后的DataFrame与预期一致：

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