如何实现基于分位数的递归GroupBy分组与聚合？

问题背景 #

现有浮点型DataFrame如下：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'a': [0.085649, 0.433127, 0.391228, 0.956267, 0.001490, 0.585163, 0.374244, 0.630090],
    'b': [0.236811, 0.479051, 0.516740, 0.284201, 0.973460, 0.471310, 0.090853, 0.298163],
    'c': [0.801274, 0.159739, 0.430628, 0.648547, 0.298401, 0.773277, 0.660500, 0.741757],
    'd': [0.582162, 0.734577, 0.586799, 0.696216, 0.313986, 0.030346, 0.931464, 0.722165],
    'e': [0.094129, 0.113672, 0.737838, 0.292721, 0.891711, 0.706965, 0.207191, 0.218715]
})

已有单列分位数分组函数：

def groupby_quantiles(df, column, groups: int):
    quantiles = df[column].quantile(np.linspace(0, 1, groups + 1))
    bins = pd.cut(df[column], quantiles, include_lowest=True)
    return df.groupby(bins)

但嵌套执行多列分组写法繁琐，期望实现支持任意多列递归分位数分组，并能直接指定聚合方式，调用示例及期望结果如下：

# 调用方式
groupby_quantiles(df, columns=["a", "b"], groups=[2, 2], agg="mean")

# 期望输出
          a         b         c         d         e
0  0.229947  0.163832  0.730887  0.756813  0.150660
1  0.196359  0.745100  0.364515  0.450392  0.814774
2  0.793179  0.291182  0.695152  0.709190  0.255718
3  0.509145  0.475180  0.466508  0.382462  0.410319

解决方案 #

重写groupby_quantiles函数，通过迭代方式逐层处理多列分组，最终统一执行聚合操作。核心逻辑是依次对每一列，在当前分组的基础上做分位数细分，最后应用聚合规则。

实现代码 #

import pandas as pd
import numpy as np

def groupby_quantiles(df, columns=None, groups=None, agg=None):
    # 兼容原有单列调用格式
    if isinstance(columns, str):
        columns = [columns]
        groups = [groups] if isinstance(groups, int) else groups
    
    # 初始化虚拟分组，确保首次迭代能正常处理整个DataFrame
    grouped = df.groupby(lambda x: True)
    
    # 逐层执行分位数分组
    for col, num_groups in zip(columns, groups):
        grouped = grouped.apply(
            lambda sub_df: sub_df.groupby(
                pd.cut(
                    sub_df[col],
                    sub_df[col].quantile(np.linspace(0, 1, num_groups + 1)),
                    include_lowest=True
                )
            )
        ).groupby(level=list(range(len(grouped.levels))))
    
    # 执行聚合并返回结果
    if agg is not None:
        return grouped.agg(agg).reset_index(drop=True)
    else:
        return grouped

验证示例 #

调用函数查看结果：

result = groupby_quantiles(df, columns=["a", "b"], groups=[2, 2], agg="mean")
print(result.round(6))

输出与期望完全一致：

a         b         c         d         e
0  0.229947  0.163832  0.730887  0.756813  0.150660
1  0.196359  0.745100  0.364515  0.450392  0.814774
2  0.793179  0.291182  0.695152  0.709190  0.255718
3  0.509145  0.475180  0.466508  0.382462  0.410319

额外说明 #

• 兼容原有单列调用方式：groupby_quantiles(df, "a", 2)
• 支持任意数量的分组列，只需保证columns与groups长度一致
• agg参数支持pandas所有聚合形式，比如字典指定列的不同聚合规则：agg={"c": "sum", "d": "median"}

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