多列分组DataFrame转嵌套字典的高效实现方案咨询

问题背景 #

现有一个包含10列的DataFrame，结构示例如下：

df:
   |  x |  y |  z |  t |     a    |  b                    | c  | ....
1: | x1 | y1 | z1 | t1 | [a1, a2] | {b1: 1, b2: 2}        | 0  | ....
2: | x2 | y2 | z2 | t2 | [a3, a4] | {b1: 3, b2: 4}        | 2  | ....
3: | x1 | y3 | z2 | t1 | [a1, a4] | {b1: 1, b2: 3}        | 2  | ....
4: | x3 | y1 | z5 | t3 | [a4, a5] | {b1: 6, b2: 2, b3: 1} | 24 | ....
                    .
                    .
                    .

需要按x、y、z、t四列的取值拆分DataFrame，转换为如下嵌套字典结构：

x1:
    y1:
         z1:
             t1: df
             t2: df
         z2:
             t1: df
             t2: df
         z3:
             t1: df
             t2: df
    y2:
         z1:
             t1: df
             t2: df
         z2:
             t1: df
             t2: df
         z3:
             t1: df
             t2: df
x2:
    y1:
         z1:
             t1: df
             t2: df
         z2:
             t1: df
             t2: df
         z3:
             t1: df
             t2: df
.
.
.

原实现代码通过多层循环嵌套调用groupby，虽然能生成目标结构，但处理大型DataFrame时计算开销极高：

grouped_df = dict(list(df.groupby('x')))
for x_elem in [*grouped_df]:
    grouped_df[x_elem] = dict(list(grouped_df[x_elem].groupby('y')))
    for y_elem in [*grouped_df[x_elem]]:
        grouped_df[x_elem][y_elem] = dict(list(grouped_df[x_elem][y_elem].groupby('z')))
        for z_elem in [*grouped_df[x_elem][y_elem]]:
            grouped_df[x_elem][y_elem][z_elem] = dict(list(grouped_df[x_elem][y_elem][z_elem].groupby('t')))

高效优化方案 #

原代码的核心问题是多次嵌套调用groupby，每次分组都会对子集重新计算，带来大量重复开销。以下是两种高效实现方式：

方案一：一次性多列分组 + 递归构建嵌套字典 #

利用pandas的多列分组特性，仅做一次分组操作，再通过递归将多级分组结果转换为嵌套字典：

def build_nested_dict(grouped_obj):
    # 当分组只剩最后一级时，直接转换为字典
    if len(grouped_obj.grouper.names) == 1:
        return dict(grouped_obj)
    # 递归处理每一层分组
    return {key: build_nested_dict(group) for key, group in grouped_obj}

# 一次性按x/y/z/t四列分组，再递归构建嵌套字典
nested_result = build_nested_dict(df.groupby(['x', 'y', 'z', 't']))

方案二：利用多级分组结果直接推导（非递归） #

如果偏好非递归实现，可以直接遍历多级分组的结果，从内层到外层逐层构建字典：

# 一次性完成四列分组
grouped = df.groupby(['x', 'y', 'z', 't'])

# 从最内层t开始构建，逐层向外嵌套
t_level = {t_val: sub_df for (x_val, y_val, z_val, t_val), sub_df in grouped}
z_level = {z_val: {k: v for k, v in t_level.items() if k[2] == z_val} for z_val in df['z'].unique()}
y_level = {y_val: {k: v for k, v in z_level.items() if any(t[1] == y_val for t in list(v.keys())[0])} for y_val in df['y'].unique()}
final_result = {x_val: {k: v for k, v in y_level.items() if any(t[0] == x_val for t in list(v.keys())[0])} for x_val in df['x'].unique()}

优化原理 #

• 原方案需要执行1 + Nx + Nx*Ny + Nx*Ny*Nz次groupby操作（Nx/Ny/Nz分别为x/y/z列的唯一值数量），时间复杂度极高。
• 优化后的方案仅执行1次groupby，之后仅通过遍历分组结果构建字典，时间复杂度从O(M*K)降至O(M)（M为DataFrame行数，K为分组层级数），在大型数据集下性能提升显著。

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