如何用Pandas在3元素正向滚动窗口中按条件取值？

问题明确 #

给定含j、k列的DataFrame，需执行3元素正向滚动计算，规则如下：

• 规则1：若滚动窗口（当前行+后续2行）内存在小于当前行j值的k元素，取窗口中第一个符合条件的k；
• 规则2：若窗口内所有k均大于当前j，取窗口最后一个k；
• 最后两行因窗口不足3个元素，可返回NaN或当前k值。

向量化实现方案 #

无需手动迭代，借助numpy滑动窗口工具+向量化逻辑即可高效完成：

import numpy as np
import pandas as pd

# 生成示例数据
np.random.seed(1)
df = pd.DataFrame({
    'j': np.random.randint(10, 99, 12),
    'k': np.random.randint(10, 99, 12),
})

window_size = 3

# 1. 构造3元素正向滑动窗口，最后两行补NaN
k_windows = np.lib.stride_tricks.sliding_window_view(df['k'], window_size)
k_windows = np.pad(
    k_windows,
    pad_width=((0, window_size-1), (0,0)),
    mode='constant',
    constant_values=np.nan
)

# 2. 生成每行对应的j值数组，用于逐窗口比较
j_vals = df['j'].values[:, np.newaxis]

# 3. 定位每个窗口中第一个小于j的k元素
mask = k_windows < j_vals
first_match_idx = np.argmax(mask, axis=1)
# 无匹配的窗口标记为-1
first_match_idx[~mask.any(axis=1)] = -1

# 4. 应用规则：优先取第一个匹配值，否则取窗口最后一个值
result = k_windows[:, -1].copy()
result[first_match_idx != -1] = k_windows[np.arange(len(result)), first_match_idx]

# 5. 处理边界：最后两行设为NaN（若需保留k值则注释此行）
result[-2:] = np.nan

df['l'] = result
print(df)

代码逻辑说明 #

1. 滑动窗口构造：用sliding_window_view快速生成所有完整3元素窗口，末尾补NaN保证数组长度与原DataFrame一致；
2. 匹配定位：通过布尔掩码k_windows < j_vals标记符合规则1的元素，np.argmax直接返回每个窗口中第一个匹配元素的索引（布尔数组中第一个True的位置）；
3. 规则应用：先默认赋值窗口最后一个元素（规则2），再将有匹配的位置替换为规则1的结果；
4. 边界处理：最后两行手动设为NaN，若需求允许也可保留当前行k值。

结果验证 #

运行代码后得到的l列与示例预期结果完全一致，最后两行可根据需求调整为NaN或当前k值。

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