Python按ID分组应用pivotid函数识别局部高低点遇阻求助

你遇到的核心问题是原pivotid函数依赖全局DataFrame的索引和数据，分组后无法适配组内的局部上下文。下面提供两种可行的重构方案，解决分组后应用的问题：

方案1：重构函数处理单个分组DataFrame #

将原函数修改为接收整个分组的DataFrame，在组内逐行计算局部高低点，避免依赖全局索引和数据：

def pivotid_group(df_group, n1, n2):
    df = df_group.copy()
    group_length = len(df)
    df['pivot'] = 0  # 初始化pivot列
    
    for idx in range(group_length):
        # 跳过组内前后不足n1/n2周期的边界行
        if idx < n1 or idx >= group_length - n2:
            continue
        
        # 获取当前行前后n1到n2的窗口数据
        window_low = df['Low'].iloc[idx-n1:idx+n2+1]
        window_high = df['High'].iloc[idx-n1:idx+n2+1]
        
        # 判断是否为局部高低点
        is_local_low = (window_low.min() == df['Low'].iloc[idx])
        is_local_high = (window_high.max() == df['High'].iloc[idx])
        
        # 映射到对应的pivot值
        if is_local_low and is_local_high:
            df['pivot'].iloc[idx] = 3
        elif is_local_low:
            df['pivot'].iloc[idx] = 1
        elif is_local_high:
            df['pivot'].iloc[idx] = 2
    return df

调用方式 #

假设你的分组列名为ID（如存储EURUSD Curncy、XAUUSD Curncy的列），直接分组后应用函数：

result_df = df.groupby('ID', group_keys=False).apply(pivotid_group, n1=2, n2=2)

方案2：用Rolling窗口+向量化计算优化效率 #

如果你的数据量较大，逐行循环效率较低，可以用Pandas的rolling窗口结合向量化操作实现，避免循环：

def get_pivot_optimized(df_group, n1, n2):
    df = df_group.copy()
    window_size = n1 + n2 + 1  # 前后n1/n2周期的总窗口大小
    
    # 计算每个点在居中窗口内是否为最低/最高点
    df['is_low'] = df['Low'] == df['Low'].rolling(window=window_size, center=True).min()
    df['is_high'] = df['High'] == df['High'].rolling(window=window_size, center=True).max()
    
    # 手动修正边界行：前后不足n1/n2的位置不参与判断
    df.loc[:n1-1, ['is_low', 'is_high']] = False
    df.loc[len(df)-n2:, ['is_low', 'is_high']] = False
    
    # 映射pivot值
    df['pivot'] = 0
    df.loc[df['is_low'] & df['is_high'], 'pivot'] = 3
    df.loc[df['is_low'] & ~df['is_high'], 'pivot'] = 1
    df.loc[~df['is_low'] & df['is_high'], 'pivot'] = 2
    
    # 可选：删除临时列
    df = df.drop(['is_low', 'is_high'], axis=1)
    return df

调用方式 #

result_df = df.groupby('ID', group_keys=False).apply(get_pivot_optimized, n1=2, n2=2)

关键注意事项 #

• 分组前排序：确保每个分组内的时间序列是按时间升序排列的，否则窗口计算会出错。
• 索引问题：原函数依赖全局索引的判断（如idx < n1），分组后必须改用组内的相对位置，这是导致原方案返回NaN的核心原因。
• group_keys=False：避免结果出现多级索引，保持原DataFrame的结构简洁。

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