如何将Pandas DataFrame转换为含flag字段的嵌套字典

初始DataFrame定义 #

import pandas as pd

my_list = [
    ['Japan', 'Flowers'],
    ['United States', 'Guns', 'yes'],
    ['Japan', 'Sushi'],
    ['South Korea', 'Sunscreen']
]

df = pd.DataFrame(my_list, columns=["country", "sector", "flag"])

现有代码及问题 #

现有分组循环代码未提取flag字段，需求是当flag不为None时，将其添加到对应结构中（例如美国的Guns条目需包含"flag": "yes"）。

现有代码 #

out = []
for idx, g in df.groupby("country"):
    out.append({"name": idx})
    ids = {}
    for i, s in g["sector"].iteritems():
        ids.setdefault(s, []).append(i)
    out[-1]["groups"] = [{"name": k, "ids": v} for k, v in ids.items()]


out = {"groups": out}
print(out)

现有生成结果 #

{
    "groups": [
        {
            "name": "Japan",
            "groups": [
                {"name": "Flowers", "ids": [0]},
                {"name": "Sushi", "ids": [2]}
            ]
        },
        {"name": "South Korea", "groups": [{"name": "Sunscreen", "ids": [3]}]},
        {"name": "United States", "groups": [{"name": "Guns", "ids": [1]}]}
    ]
}

修改后的代码 #

调整循环逻辑，同时获取sector和flag字段，在构建子组字典时判断flag是否非空，若不为空则添加该字段：

import pandas as pd

my_list = [
    ['Japan', 'Flowers'],
    ['United States', 'Guns', 'yes'],
    ['Japan', 'Sushi'],
    ['South Korea', 'Sunscreen']
]

df = pd.DataFrame(my_list, columns=["country", "sector", "flag"])

out = []
for idx, g in df.groupby("country"):
    out.append({"name": idx})
    sector_data = {}
    # 同时遍历索引、sector和flag
    for i, row in g.iterrows():
        s = row["sector"]
        flag_val = row["flag"]
        # 初始化sector对应的条目：ids列表和flag值
        if s not in sector_data:
            sector_data[s] = {"ids": [], "flag": None}
        sector_data[s]["ids"].append(i)
        # 仅当flag不为空时更新（避免覆盖None）
        if pd.notna(flag_val):
            sector_data[s]["flag"] = flag_val
    # 构建groups列表，过滤掉flag为None的字段
    out[-1]["groups"] = []
    for k, v in sector_data.items():
        group_item = {"name": k, "ids": v["ids"]}
        if v["flag"] is not None:
            group_item["flag"] = v["flag"]
        out[-1]["groups"].append(group_item)

out = {"groups": out}
print(out)

修改后生成的结果 #

{
    "groups": [
        {
            "name": "Japan",
            "groups": [
                {"name": "Flowers", "ids": [0]},
                {"name": "Sushi", "ids": [2]}
            ]
        },
        {"name": "South Korea", "groups": [{"name": "Sunscreen", "ids": [3]}]},
        {"name": "United States", "groups": [{"name": "Guns", "ids": [1], "flag": "yes"}]}
    ]
}

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