Pandas GroupBy后字符串列表无法比对数值求缺失值如何修复

需求说明 #

现有数据框初始结构如下图左侧所示，需按Name字段分组，对比基准数值列表[1,2,3,4,5]，统计每个分组缺失的数字，期望输出效果如下图右侧所示。

原代码运行后，GroupBy聚合得到的Number of stars列为字符串组成的列表，数值比对、集合差集运算全部出错，所有分组的差集都返回了完整的基准集合。

问题根因 #

读取CSV后，Number of stars列的每一行值是类似"1,3,2,1,2"的整段逗号分隔字符串，原代码直接对该列做分组聚合收集列表，没有做字符串拆分和整数类型转换，后续转集合得到的是字符串集合，和基准的整数集合类型不匹配，差集运算时判定两类元素完全不重合，因此返回错误结果。

修复要点 #

• 先拆分每行的逗号分隔字符串，将拆分后的单个值转为整数，再做分组聚合
• 聚合时合并同组所有行的整数为一个列表，保证参与集合运算的元素都是整数类型
• 集合差集运算时保证两边元素类型一致，避免类型不匹配导致的计算错误

修复后完整代码 #

import pandas as pd
from io import StringIO

csvfile = StringIO("""
Name    Number of stars
Benjamin    1,3,2,1,2
Benjamin    2,5,1,3
Emma    2,1,1,4,4,2
Ethan   2,5,4
Emma    2,2,2
Ethan   5,4,4,1,1,1
Olivia  4,1,3,5""")

df = pd.read_csv(csvfile, sep = '\t', engine='python')

# 拆分字符串转整数后再分组聚合
df_1 = df.assign(
    star_list = df['Number of stars'].str.split(',').map(lambda x: [int(num) for num in x])
).groupby('Name')['star_list'].agg(lambda seq: [n for item in seq for n in item]).reset_index()

base_set = {1,2,3,4,5}
df_1['all stars'] = [sorted(base_set) for _ in range(len(df_1))]
# 计算缺失值
df_1['missing'] = df_1['star_list'].map(lambda x: sorted(base_set - set(x)))

print(df_1)

正确运行输出 #

Name                                 star_list        all stars missing
0  Benjamin        [1, 3, 2, 1, 2, 2, 5, 1, 3]  [1, 2, 3, 4, 5]     [4]
1      Emma  [2, 1, 1, 4, 4, 2, 2, 2, 2]  [1, 2, 3, 4, 5]  [3, 5]
2     Ethan     [2, 5, 4, 5, 4, 4, 1, 1, 1]  [1, 2, 3, 4, 5]     [3]
3    Olivia                    [4, 1, 3, 5]  [1, 2, 3, 4, 5]     [2]

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