Pandas转嵌套字典技巧分享

在数据处理中，我们经常需要将结构化的表格数据（如Pandas DataFrame）转换为更灵活的嵌套字典格式，以便于后续的数据查询或业务逻辑处理。一个常见的需求是，将DataFrame中的两列作为嵌套字典的键，而其余列的值则聚合为一个列表作为最内层字典的值。

原始实现与局限

考虑一个包含公司（Company）、产品（Product）以及相关生产数据（Production Cost, Development Time, Launch Year）的DataFrame。我们的目标是创建一个字典，其中外层键是公司名，内层键是产品名，对应的值是生产数据的列表。

以下是使用df.iterrows()方法实现这一目标的常见代码：

import pandas as pd

# 示例DataFrame，通常从Excel文件读取
df = pd.DataFrame({
    "Company": ["TechCorp", "Innovate Inc", "Green Solutions", "Future Dynamics"],
    "Product": ["TC100", "IN200", "GS300", "FD400"],
    "Production Cost": [10000, 15000, 12000, 18000],
    "Development Time": [6, 9, 8, 12],
    "Launch Year": [2023, 2024, 2023, 2025]
})

nested_dict = {}
for index, row in df.iterrows():
    company = row['Company']
    product = row['Product']
    # 提取其他列的值并转换为列表
    values = row[['Production Cost', 'Development Time', 'Launch Year']].tolist()

    if company not in nested_dict:
        nested_dict[company] = {}
    nested_dict[company][product] = values

print(nested_dict)

这段代码能够正确生成所需的嵌套字典结构：

{'TechCorp': {'TC100': [10000, 6, 2023]}, 'Innovate Inc': {'IN200': [15000, 9, 2024]}, 'Green Solutions': {'GS300': [12000, 8, 2023]}, 'Future Dynamics': {'FD400': [18000, 12, 2025]}}

然而，df.iterrows()在处理大型DataFrame时效率相对较低。每次迭代都会返回一个Series对象，这涉及到额外的开销。同时，手动检查字典中是否存在外层键 (if company not in nested_dict:) 增加了代码的冗余。

优化方案：Defaultdict与df.values的结合

为了提高效率并简化代码，我们可以利用Python标准库collections中的defaultdict和Pandas DataFrame的底层NumPy数组表示df.values。

1. collections.defaultdict: defaultdict是dict的一个子类，它允许在访问不存在的键时自动创建该键并赋予一个默认值。在本例中，我们可以将defaultdict(dict)作为外层字典，这样当遇到新的公司名时，会自动创建一个空的内层字典。
2. df.values与扩展解包: df.values属性返回DataFrame中所有数据的NumPy数组表示。直接遍历这个数组比iterrows()更高效。通过使用Python的扩展解包（*操作符），我们可以方便地将一行中的前几个元素分配给特定变量，而将剩余的所有元素收集到一个列表中。

以下是优化后的代码实现：

from collections import defaultdict
import pandas as pd

# 示例DataFrame
df = pd.DataFrame({
    "Company": ["TechCorp", "Innovate Inc", "Green Solutions", "Future Dynamics"],
    "Product": ["TC100", "IN200", "GS300", "FD400"],
    "Production Cost": [10000, 15000, 12000, 18000],
    "Development Time": [6, 9, 8, 12],
    "Launch Year": [2023, 2024, 2023, 2025]
})

# 为了演示defaultdict的优势，我们添加一个TechCorp的新产品
df.loc[len(df)] = ['TechCorp', 'TC200', 20000, 12, 2025]

nested_dict = defaultdict(dict)
# 遍历df.values，使用扩展解包
for company, product, *values in df.values:
    nested_dict[company][product] = values

print(dict(nested_dict)) # 将defaultdict转换为普通dict以便输出

代码解析：

• nested_dict = defaultdict(dict)：初始化一个defaultdict，其默认工厂函数是dict。这意味着如果nested_dict中尝试访问一个不存在的company键，它会自动创建一个空的普通字典作为该键的值。
• for company, product, *values in df.values:：
  • df.values将DataFrame转换为一个NumPy数组。
  • 每次迭代获取数组中的一行。
  • company和product分别捕获行的前两个元素。
  • *values是Python的扩展解包语法，它会收集行中剩余的所有元素，并将它们作为一个列表赋值给values变量。这完美匹配了我们的需求，即其余列的值作为一个列表。

优化后的输出（包含新增数据）：

{
 'TechCorp': {'TC100': [10000, 6, 2023], 'TC200': [20000, 12, 2025]},
 'Innovate Inc': {'IN200': [15000, 9, 2024]},
 'Green Solutions': {'GS300': [12000, 8, 2023]},
 'Future Dynamics': {'FD400': [18000, 12, 2025]}
}

可以看到，TechCorp下的两个产品及其数据都被正确地聚合到了同一个外层字典中，证明了defaultdict的有效性。

注意事项

• 列顺序的重要性: 使用df.values时，数据的顺序至关重要。确保DataFrame中作为键的列（Company, Product）位于需要作为值列表的列之前。如果列顺序不符合预期，需要在使用df.values前对DataFrame进行列重排。
• 性能提升: 对于大型数据集，直接遍历df.values通常比iterrows()快得多，因为df.values操作直接在NumPy数组上进行，避免了Pandas Series对象的创建开销。
• 代码简洁性: defaultdict消除了手动检查键是否存在的if语句，使得代码更加简洁和易读。扩展解包也极大地简化了值列表的创建过程。
• 数据类型: df.values返回的NumPy数组中的元素会保留其原始数据类型。在大多数情况下，这不会有问题，但如果需要特定类型转换，可能需要在*values捕获后进行处理。

总结

通过巧妙地结合collections.defaultdict和df.values的扩展解包功能，我们可以显著优化从Pandas DataFrame构建复杂嵌套字典的过程。这种方法不仅提升了代码的执行效率，尤其是在处理大规模数据时，还增强了代码的简洁性和可读性，是Python数据处理中值得推荐的实践。