如何结合groupby与numpy函数实现个股连续下跌状态判断？

首先，你的原始代码存在两个核心问题：

1. 循环从i=0开始，当i<2时访问close[i-2]会触发索引越界错误
2. 使用groupby.apply()时，函数接收的是整个分组DataFrame而非单独的close列，导致KeyError

方案1：修正原函数并适配分组逻辑 #

先修复函数的索引问题，再通过groupby['close'].transform()传入每个分组的close序列（transform会自动对齐结果索引）：

import numpy as np
import pandas as pd

def is_downtrending(close):
    out = np.full(close.shape, False)
    # 从i=2开始遍历，确保i-2是有效索引
    for i in range(2, close.shape[0]):
        if (close[i] < close[i-1]) and (close[i-1] < close[i-2]):
            out[i] = True
    return out

# 使用transform自动对齐结果到原DataFrame
df['is_downtrending'] = df.groupby('stock_id')['close'].transform(is_downtrending)

方案2：矢量化操作（高效处理大数据集） #

循环在22万行数据上效率低下，推荐用pandas内置方法实现相同逻辑，无需自定义函数：

方法A：利用diff和shift

# 连续两日下跌等价于：当日<前日 且 前日<前日的前日
df['is_downtrending'] = df.groupby('stock_id')['close'].transform(
    lambda x: (x.diff() < 0) & (x.shift(1).diff() < 0)
).fillna(False)

方法B：滚动窗口判断

# 滚动窗口检查连续3天是否呈递减趋势（旧→新：a > b > c）
df['is_downtrending'] = df.groupby('stock_id')['close'].rolling(3).apply(
    lambda win: win[0] > win[1] > win[2], raw=True
).reset_index(level=0, drop=True).fillna(False)

关键说明 #

• groupby['close'].transform()会将每个分组的close列作为Series传入函数，完美匹配原函数的输入要求，且自动对齐结果索引
• 矢量化操作比循环快数倍，更适合处理大规模数据集
• fillna(False)是为了处理每个分组前两行无法判断的情况（保持默认False）

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