如何在DataFrame中按相同区域合并列中同键字典并聚合数据

原始数据 #

import pandas as pd

df = pd.DataFrame({
    'ID': ['a001', 'a002', 'a003'],
    'area': ['NY', 'SF', 'NY'],
    'data_info': [
        [{'color': 'Yellow', 'count': 3, 'weight': 5}, {'color': 'Blue', 'count': 2, 'weight': 11}, {'color': 'Red', 'count': 7, 'weight': 3}],
        [{'color': 'Green', 'count': 1, 'weight': 14}, {'color': 'Yellow', 'count': 9, 'weight': 2}],
        [{'color': 'Blue', 'count': 5, 'weight': 6}, {'color': 'Black', 'count': 2, 'weight': 15}]
    ]
})

需求 #

将相同area的行中data_info列的嵌套字典合并，按color键对count和weight进行求和聚合。

解决代码 #

def aggregate_data(group):
    # 扁平化分组内所有data_info的字典列表
    all_items = [item for sublist in group['data_info'] for item in sublist]
    # 按color聚合统计值
    agg_result = {}
    for item in all_items:
        color = item['color']
        if color not in agg_result:
            agg_result[color] = {'color': color, 'count': 0, 'weight': 0}
        agg_result[color]['count'] += item['count']
        agg_result[color]['weight'] += item['weight']
    # 转换为列表返回
    return list(agg_result.values())

# 分组聚合
result_df = df.groupby('area').apply(aggregate_data).reset_index(name='data_info')

结果验证 #

执行代码后，result_df的输出如下：

area                                          data_info
0   NY  [{'color': 'Yellow', 'count': 3, 'weight': 5}, {'color': 'Blue', 'count': 7, 'weight': 17}, {'color': 'Red', 'count': 7, 'weight': 3}, {'color': 'Black', 'count': 2, 'weight': 15}]
1   SF  [{'color': 'Green', 'count': 1, 'weight': 14}, {'color': 'Yellow', 'count': 9, 'weight': 2}]

其中NY区域的Blue对应的count为2+5=7，weight为11+6=17，完全符合预期。

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