如何在Pandas中为多股票DataFrame添加连续收阴判断布尔列

问题背景 #

你有一个包含多只股票的DataFrame，结构如下：

index        date symbol  stock_id  open  high     low  close    volume      vwap
0            0  2021-10-11    BVN        13  7.69  7.98  7.5600   7.61    879710  7.782174
1            1  2021-10-12    BVN        13  7.67  8.08  7.5803   8.02    794436  7.967061
2            2  2021-10-13    BVN        13  8.12  8.36  8.0900   8.16    716012  8.231286
3            3  2021-10-14    BVN        13  8.26  8.29  8.0500   8.28    586091  8.185899
4            4  2021-10-15    BVN        13  8.18  8.44  8.0600   8.44   1278409  8.284539
...        ...         ...    ...       ...   ...   ...     ...    ...       ...       ...
227774  227774  2022-10-04   ERIC     11000  6.27  6.32  6.2400   6.29  14655189  6.280157
227775  227775  2022-10-05   ERIC     11000  6.17  6.31  6.1500   6.29  10569193  6.219965
227776  227776  2022-10-06   ERIC     11000  6.20  6.25  6.1800   6.22   7918812  6.217198
227777  227777  2022-10-07   ERIC     11000  6.17  6.19  6.0800   6.10   9671252  6.135976
227778  227778  2022-10-10   ERIC     11000  6.13  6.15  6.0200   6.04   6310661  6.066256

[227779 rows x 10 columns]

需要为每只股票新增布尔列is_downtrending，当该股票**连续两日收盘价低于前一日（即当日收盘价<前一日，且前一日收盘价<前两日）**时，该列值为True，否则为False。预期效果示例如下：

index        date symbol  stock_id  open  high     low  close    volume      vwap  is_downtrending
0            0  2021-10-11    BVN        13  7.69  7.98  7.5600   7.61    879710  7.782174  False
1            1  2021-10-12    BVN        13  7.67  8.08  7.5803   8.02    794436  7.967061  False
2            2  2021-10-13    BVN        13  8.12  8.36  8.0900   8.16    716012  8.231286  False
3            3  2021-10-14    BVN        13  8.26  8.29  8.0500   8.28    586091  8.185899  True
4            4  2021-10-15    BVN        13  8.18  8.44  8.0600   8.44   1278409  8.284539  True
...        ...         ...    ...       ...   ...   ...     ...    ...       ...       ...
227774  227774  2022-10-04   ERIC     11000  6.27  6.32  6.2400   6.29  14655189  6.280157  True
227775  227775  2022-10-05   ERIC     11000  6.17  6.31  6.1500   6.29  10569193  6.219965  True
227776  227776  2022-10-06   ERIC     11000  6.20  6.25  6.1800   6.22   7918812  6.217198  True
227777  227777  2022-10-07   ERIC     11000  6.17  6.19  6.0800   6.10   9671252  6.135976  False
227778  227778  2022-10-10   ERIC     11000  6.13  6.15  6.0200   6.04   6310661  6.066256  True

[227779 rows x 11 columns]

原代码的问题 #

你尝试的代码存在以下问题：

1. 索引越界：循环从i=0开始，当i=0或i=1时，i-1、i-2会取到负数索引，导致错误的取值（会取到列的最后一个元素）。
2. 结果拼接错误：每次循环都直接赋值final_df = ohlc，最终只会保留最后一只股票的数据，没有累加所有分组结果。
3. 未处理边界情况：每组的前两行没有足够的历史数据，应该默认设为False。

正确解法 #

方法一：修复原代码逻辑 #

针对原代码的问题进行修正：

import pandas as pd
import numpy as np

ticker_group = df.groupby(['stock_id'])
final_df = pd.DataFrame()

def is_downtrending(close):
    out = np.full(close.shape[0], False)
    # 从第3条数据开始（索引2），确保有前两日的数据
    for i in range(2, close.shape[0]):
        # 判断当日收盘价<前一日，且前一日<前两日
        if (close.iloc[i] < close.iloc[i-1]) and (close.iloc[i-1] < close.iloc[i-2]):
            out[i] = True
    return out

for index, ohlc in ticker_group:
    ohlc['is_downtrending'] = is_downtrending(ohlc['close'])
    # 用concat累加所有分组数据
    final_df = pd.concat([final_df, ohlc], ignore_index=True)

方法二：Pandas矢量化操作（更高效） #

避免循环，用Pandas的groupby和shift方法实现矢量化计算，适合处理大规模数据集：

# 1. 按股票分组，标记每日是否收阴（收盘价<前一日）
df['is_red'] = df.groupby('stock_id')['close'].apply(lambda x: x < x.shift(1))

# 2. 判断是否连续两日收阴：当日收阴且前一日也收阴
df['is_downtrending'] = df['is_red'] & df['is_red'].shift(1)

# 3. 填充缺失值（每组前两行无足够数据，设为False）
df['is_downtrending'] = df['is_downtrending'].fillna(False)

如果你的实际需求是标记连续下跌趋势中的所有日期（如示例中BVN的第3、4行，即使第4行收盘价上涨仍标记为True），可以调整为以下逻辑：

# 先标记首次出现连续两日收阴的点
df['consec_red'] = df.groupby('stock_id')['is_red'].rolling(2).sum().reset_index(level=0, drop=True) == 2

# 用cummax将标记点后的所有日期设为True，直到趋势结束
df['is_downtrending'] = df.groupby('stock_id')['consec_red'].transform(lambda x: x.cummax())

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