如何无循环实现DataFrame指定前后行的极值判断？

给定如下结构的pandas DataFrame：

import pandas as pd

df = pd.DataFrame(
    [
        [5, 8],
        [8, 10],
        [3, 15],
        [16, 20],
        [12, 21],
        [5, 9],
        [10, 12],
        [20, 22],
        [4, 10],
        [7, 13],
        [9, 15],
        [6, 9],
    ],
    columns=list("lh"),
)

需求是为每行判断：在指定的前n1行、后n2行范围内（示例中n1=2，n2=2），该行是否为l列的最小值或h列的最大值，以此生成pivot列：

• 若同时是l列最小值和h列最大值，返回3
• 仅为l列最小值，返回1
• 仅为h列最大值，返回2
• 都不是，返回0
• 边界行（窗口超出DataFrame范围的行）返回0

原实现使用循环函数结合apply，处理大数据量时效率极低，现寻求无循环的pandas高效实现方式。

原循环实现代码：

def pivotid(df1, l, n1, n2):  # n1 n2 before and after candle L
    if l - n1 < 0 or l + n2 >= len(df1):
        return 0

    pividlow = 1
    pividhigh = 1
    for i in range(l - n1, l + n2 + 1):
        if df1.l[l] > df1.l[i]:
            pividlow = 0
        if df1.h[l] < df1.h[i]:
            pividhigh = 0

    if pividlow and pividhigh:
        return 3
    elif pividlow:
        return 1
    elif pividhigh:
        return 2
    else:
        return 0

df['pivot'] = df.apply(lambda x: pivotid(df, x.name, 2, 2), axis=1)

预期结果：

l   h  pivot
0    5   8      0
1    8  10      0
2    3  15      1
3   16  20      0
4   12  21      2
5    5   9      1
6   10  12      0
7   20  22      2
8    4  10      1
9    7  13      0
10   9  15      0
11   6   9      0

可以利用pandas的**滑动窗口（rolling）**功能结合向量化操作实现，完全避免Python循环，大幅提升处理效率：

步骤1：定义窗口参数 #

n1 = 2
n2 = 2
window_size = n1 + n2 + 1  # 窗口总大小：前2行+当前行+后2行=5行

步骤2：计算窗口内的极值 #

使用rolling方法设置center=True，让窗口中心对齐当前行；同时用min_periods=window_size确保只有完整窗口才计算值，边界行直接返回NaN：

# 计算l列的窗口最小值
window_min_l = df['l'].rolling(window=window_size, center=True, min_periods=window_size).min()
# 计算h列的窗口最大值
window_max_h = df['h'].rolling(window=window_size, center=True, min_periods=window_size).max()

步骤3：判断每行是否为窗口极值 #

通过向量化比较生成布尔列：

is_low = (df['l'] == window_min_l)
is_high = (df['h'] == window_max_h)

步骤4：生成pivot列 #

根据布尔列的组合结果映射到对应的pivot值，边界行因NaN无法匹配，默认保持初始值0：

df['pivot'] = 0
df.loc[is_low & is_high, 'pivot'] = 3
df.loc[is_low & ~is_high, 'pivot'] = 1
df.loc[~is_low & is_high, 'pivot'] = 2

完整代码 #

import pandas as pd

df = pd.DataFrame(
    [
        [5, 8],
        [8, 10],
        [3, 15],
        [16, 20],
        [12, 21],
        [5, 9],
        [10, 12],
        [20, 22],
        [4, 10],
        [7, 13],
        [9, 15],
        [6, 9],
    ],
    columns=list("lh"),
)

n1 = 2
n2 = 2
window_size = n1 + n2 + 1

# 计算窗口极值
window_min_l = df['l'].rolling(window=window_size, center=True, min_periods=window_size).min()
window_max_h = df['h'].rolling(window=window_size, center=True, min_periods=window_size).max()

# 判断极值
is_low = (df['l'] == window_min_l)
is_high = (df['h'] == window_max_h)

# 生成pivot列
df['pivot'] = 0
df.loc[is_low & is_high, 'pivot'] = 3
df.loc[is_low & ~is_high, 'pivot'] = 1
df.loc[~is_low & is_high, 'pivot'] = 2

print(df)

效果说明 #

运行后输出与预期结果完全一致，且该方案基于pandas底层C实现的滑动窗口操作，处理大数据量时效率比原循环实现高几个数量级。

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