求助：如何在多列及复杂函数中使用高效Pandas Groupby类

针对你遇到的两个问题，直接修改原Groupby类的核心逻辑即可解决，以下是具体方案：

1. 多列分组支持修复 #

原类的初始化逻辑仅处理单列分组的情况，需调整为接受列名列表，并通过多列组合的唯一标识生成分组映射：

import pandas as pd
import numpy as np
from collections import defaultdict

class FastGroupby:
    def __init__(self, df, group_cols):
        self.df = df.copy()
        # 兼容单列/多列分组参数，统一转为列表
        self.group_cols = group_cols if isinstance(group_cols, list) else [group_cols]
        # 生成多列组合的唯一分组键（用tuple保证可哈希）
        self.df['_temp_group_key'] = self.df[self.group_cols].apply(tuple, axis=1)
        self.groups = defaultdict(list)
        # 构建分组索引映射
        for idx, key in enumerate(self.df['_temp_group_key']):
            self.groups[key].append(idx)
        # 清理临时列
        self.df.drop(columns=['_temp_group_key'], inplace=True)

2. 多参数函数适配修改 #

原apply方法仅支持无额外参数的函数，扩展方法以支持传递任意位置参数和关键字参数：

def apply(self, func, *args, **kwargs):
        results = {}
        for group_key, indices in self.groups.items():
            group_data = self.df.iloc[indices]
            # 将额外参数传递给自定义函数
            results[group_key] = func(group_data, *args, **kwargs)
        # 将结果转换为DataFrame，自动对齐原分组列
        result_df = pd.DataFrame.from_dict(results, orient='index', columns=['result'])
        result_df.reset_index(inplace=True)
        result_df.columns = self.group_cols + ['result']
        return result_df

测试验证 #

多列分组测试

# 生成测试数据
df = pd.DataFrame({
    'A': ['foo', 'bar', 'foo', 'bar', 'foo', 'bar'],
    'B': ['one', 'one', 'two', 'two', 'one', 'one'],
    'C': np.random.randn(6),
    'D': np.random.randn(6)
})

# 多列分组初始化
fg = FastGroupby(df, group_cols=['A', 'B'])

# 自定义简单函数
def mean_c(group):
    return group['C'].mean()

# 应用函数
print(fg.apply(mean_c))

多参数复杂函数测试（以回归预测为例）

# 带多参数的自定义回归函数
def dummypredict(group, target_col, feature_col, alpha=0.1):
    from sklearn.linear_model import Ridge
    X = group[feature_col].values.reshape(-1,1)
    y = group[target_col].values
    model = Ridge(alpha=alpha)
    model.fit(X, y)
    return model.predict(X).mean()

# 传递参数应用函数
print(fg.apply(dummypredict, target_col='D', feature_col='C', alpha=0.5))

以上修改既保留了原类的高效性，又解决了多列分组和多参数函数适配的问题。

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