Python字典列表按ID分组并新增键值对的实现问题

从数据库查询得到一个字典列表，需要按id字段对数据进行分组，并基于字典中的现有数据添加past_shows、upcoming_shows等新键值对及对应计数。现有代码未达到预期效果，寻求正确实现方案。

原始数据示例 #

raw_data = [
    {"id": 1, "show_date": "2023-01-01", "status": "past"},
    {"id": 1, "show_date": "2023-06-01", "status": "upcoming"},
    {"id": 2, "show_date": "2022-12-01", "status": "past"},
    {"id": 2, "show_date": "2023-07-01", "status": "upcoming"},
    {"id": 2, "show_date": "2023-08-01", "status": "upcoming"},
]

期望输出结果 #

expected_output = [
    {
        "id": 1,
        "past_shows": 1,
        "upcoming_shows": 1,
        "shows": [
            {"show_date": "2023-01-01", "status": "past"},
            {"show_date": "2023-06-01", "status": "upcoming"}
        ]
    },
    {
        "id": 2,
        "past_shows": 1,
        "upcoming_shows": 2,
        "shows": [
            {"show_date": "2022-12-01", "status": "past"},
            {"show_date": "2023-07-01", "status": "upcoming"},
            {"show_date": "2023-08-01", "status": "upcoming"}
        ]
    }
]

尝试的代码（未达到预期） #

# 示例错误代码（常见逻辑问题）
result = []
for item in raw_data:
    existing = next((r for r in result if r["id"] == item["id"]), None)
    if not existing:
        # 错误：初始化时可能未正确设置计数初始值，或后续计数逻辑遗漏
        result.append({
            "id": item["id"],
            "past_shows": 0,
            "upcoming_shows": 0,
            "shows": []
        })
        existing = result[-1]
    existing["shows"].append(item)
    # 错误：可能未正确匹配status值，或计数逻辑触发条件错误
    if item.get("status") == "past":
        existing["past_shows"] += 1

正确实现方案 #

方法一：使用字典映射分组（高效处理无序数据） #

这种方法不需要预先排序，直接通过字典记录每个ID对应的分组信息，适合处理大规模或无序数据：

def process_shows(raw_data):
    grouped_dict = {}
    for item in raw_data:
        show_id = item["id"]
        # 初始化分组条目（如果ID未存在）
        if show_id not in grouped_dict:
            grouped_dict[show_id] = {
                "id": show_id,
                "past_shows": 0,
                "upcoming_shows": 0,
                "shows": []
            }
        # 添加当前演出到列表
        grouped_dict[show_id]["shows"].append(item)
        # 根据状态更新计数
        status = item["status"]
        if status == "past":
            grouped_dict[show_id]["past_shows"] += 1
        elif status == "upcoming":
            grouped_dict[show_id]["upcoming_shows"] += 1
    # 将字典值转换为列表返回
    return list(grouped_dict.values())

方法二：使用itertools.groupby（代码简洁，需先排序） #

groupby要求数据按分组键连续排列，因此需要先对原始数据按id排序，适合数据本身已排序或对代码简洁性要求较高的场景：

from itertools import groupby
from operator import itemgetter

def process_shows(raw_data):
    # 先按id排序，确保groupby能正确分组
    sorted_data = sorted(raw_data, key=itemgetter("id"))
    result = []
    for show_id, group in groupby(sorted_data, key=itemgetter("id")):
        shows_list = list(group)
        # 统计不同状态的演出数量
        past_count = sum(1 for show in shows_list if show["status"] == "past")
        upcoming_count = sum(1 for show in shows_list if show["status"] == "upcoming")
        # 组装结果条目
        result.append({
            "id": show_id,
            "past_shows": past_count,
            "upcoming_shows": upcoming_count,
            "shows": shows_list
        })
    return result

测试验证 #

调用上述函数并传入原始数据，即可得到与期望输出一致的结果：

raw_data = [
    {"id": 1, "show_date": "2023-01-01", "status": "past"},
    {"id": 1, "show_date": "2023-06-01", "status": "upcoming"},
    {"id": 2, "show_date": "2022-12-01", "status": "past"},
    {"id": 2, "show_date": "2023-07-01", "status": "upcoming"},
    {"id": 2, "show_date": "2023-08-01", "status": "upcoming"},
]

print(process_shows(raw_data))

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