如何高效将Pandas DataFrame转换为指定结构的嵌套字典？

我有一个包含3列的Pandas DataFrame，结构如下：

import pandas as pd

test_df = pd.DataFrame({
    'key1': [1, 1, 1, 1, 2, 2, 2],
    'key2': ['a', 'b', 'c', 'd', 'e', 'f', 'g'],
    'value': ['a-mapped', 'b-mapped', 'c-mapped', 'd-mapped', 'e-mapped', 'f-mapped', 'g-mapped']
})

输出的DataFrame如下：

key1 key2     value
0      1    a  a-mapped
1      1    b  b-mapped
2      1    c  c-mapped
3      1    d  d-mapped
4      2    e  e-mapped
5      2    f  f-mapped
6      2    g  g-mapped

我希望将其转换为嵌套字典：外层键为key1的取值，内层是key2与value组成的键值对，期望结果如下：

{
    1: {'a': 'a-mapped', 'b': 'b-mapped', 'c': 'c-mapped', 'd': 'd-mapped'},
    2: {'e': 'e-mapped', 'f': 'f-mapped', 'g': 'g-mapped'}
}

目前我通过循环实现了需求：

desired_result = {}
for key1 in test_df.key1:
    desired_result[key1] = {}
    
for idx, row in test_df.iterrows():
    desired_result[row.key1][row.key2] = row.value

请问是否存在更高效的实现方式？

回答 #

当然有，利用Pandas的内置分组和向量化操作可以大幅提升效率，尤其是在数据量较大时，避免Python层面的逐行循环：

方法1：分组后结合apply与dict #

desired_result = test_df.groupby('key1').apply(lambda x: dict(zip(x['key2'], x['value']))).to_dict()

方法2：字典推导式结合groupby #

desired_result = {group_key: dict(zip(group['key2'], group['value'])) for group_key, group in test_df.groupby('key1')}

为什么这两种方法更高效？ #

• 这两种方法依赖Pandas底层的C实现分组逻辑，比iterrows()这种逐行遍历的Python循环快得多，数据量越大，性能差距越明显。
• 原循环存在冗余操作：多次对同一个key1初始化空字典，而分组方法会直接按key1的唯一值分组，一次性生成对应内层字典。

验证结果 #

执行上述任意一种方法后，输出结果与期望完全一致：

print(desired_result)
# {1: {'a': 'a-mapped', 'b': 'b-mapped', 'c': 'c-mapped', 'd': 'd-mapped'}, 2: {'e': 'e-mapped', 'f': 'f-mapped', 'g': 'g-mapped'}}

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