如何使用groupby将列聚合为字典并生成含该字典的新列？

需求描述 #

现有包含person、year及若干flag列的DataFrame，希望保留person和year列，并生成新列flag_dict，存储对应行所有flag列的键值对字典。

示例DataFrame #

# sample dataframe
import pandas as pd
data = [["John Doe", 2018, True, False, True], ["Jane Doe", 2019, True, False, False]]
df = pd.DataFrame(data, columns=["person", "year", "flag_1", "flag_2", "flag_3"])

输出：

person  year  flag_1  flag_2  flag_3
0  John Doe  2018    True   False    True
1  Jane Doe  2019    True   False   False

用户尝试过程 #

用户先尝试重塑数据：

reshaped_df = pd.melt(
    df, id_vars=["person", "year"], value_vars=["flag_1", "flag_2", "flag_3"]
)

得到：

person  year variable  value
0  John Doe  2018   flag_1   True
1  Jane Doe  2019   flag_1   True
2  John Doe  2018   flag_2  False
3  Jane Doe  2019   flag_2  False
4  John Doe  2018   flag_3   True

再尝试生成字典：

reshaped_df.set_index(["person", "year", "variable"]).T.to_dict("list")

得到的结果是每个flag单独的键值对，不符合预期：

{('John Doe', 2018, 'flag_1'): [True],
 ('Jane Doe', 2019, 'flag_1'): [True],
 ('John Doe', 2018, 'flag_2'): [False],
 ('Jane Doe', 2019, 'flag_2'): [False],
 ('John Doe', 2018, 'flag_3'): [True],
 ('Jane Doe', 2019, 'flag_3'): [False]}

目标输出 #

person  year                                           flag_dict
0  John Doe  2018  {'flag_1': True, 'flag_2': False, 'flag_3': True}
1  Jane Doe  2019  {'flag_1': True, 'flag_2': False, 'flag_3': False}

解决方案 #

方法1：直接对每行生成字典（最简洁） #

无需groupby，直接提取flag列后对每行应用to_dict()：

# 筛选所有flag列
flag_cols = [col for col in df.columns if col.startswith('flag_')]
# 生成新列
df['flag_dict'] = df[flag_cols].apply(lambda x: x.to_dict(), axis=1)
# 保留目标列
result = df[['person', 'year', 'flag_dict']]
print(result)

方法2：使用groupby实现 #

按照person和year分组，对每组的flag列聚合为字典：

flag_cols = [col for col in df.columns if col.startswith('flag_')]

result = df.groupby(['person', 'year'], as_index=False).apply(
    lambda x: x[flag_cols].iloc[0].to_dict(),
    include_groups=False
).rename('flag_dict').reset_index()
print(result)

方法3：基于用户的melt结果改进 #

基于已有的重塑数据，分组后将每组的variable和value配对成字典：

reshaped_df = pd.melt(
    df, id_vars=["person", "year"], value_vars=["flag_1", "flag_2", "flag_3"]
)

result = reshaped_df.groupby(['person', 'year']).apply(
    lambda x: dict(zip(x['variable'], x['value'])),
    include_groups=False
).rename('flag_dict').reset_index()
print(result)

以上三种方法均能得到目标输出，其中方法1最为高效简洁。

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