如何关联Pandas DataFrame两列列表项并计算分类统计指标

现有一个包含两列列表数据的Pandas DataFrame，一列存储物品分类（cat列），另一列存储对应分类的分数（pts列），构造代码如下：

import pandas as pd

cat = [
    ['speed', 'health', 'strength', 'health'],
    ['strength', 'speed', 'speed'],
    ['strength', 'speed', 'health', 'speed']
]

pts = [
    [1, 2, 1.5, -1],
    [2, -1.5, 1.5],
    [-1, 2, 0, 1.5]
]

s_cat = pd.Series(cat, name='cat')
s_pts = pd.Series(pts, name='pts')

df = pd.concat([s_cat, s_pts], axis=1)

原始DataFrame输出：

cat              pts
0  [speed, health, strength, health]  [1, 2, 1.5, -1]
1           [strength, speed, speed]   [2, -1.5, 1.5]
2   [strength, speed, health, speed]  [-1, 2, 0, 1.5]

需求是关联两列的列表项，生成各分类的分数总和、正分数计数、负分数计数的新列，期望结果如下：

cat              pts  speed_sum  health_sum  strength_sum  speed_pos  speed_neg  health_pos  health_neg  strength_pos  strength_neg
0  [speed, health, strength, health]  [1, 2, 1.5, -1]        1.0         1.0           1.5          1          0           1           1             1             0
1           [strength, speed, speed]   [2, -1.5, 1.5]        0.0         NaN           2.0          1          1           0           0             1             0
2   [strength, speed, health, speed]  [-1, 2, 0, 1.5]        3.5         0.0          -1.0          2          0           0           0             0             1

解决方案代码 #

import pandas as pd

# 构造原始数据
cat = [
    ['speed', 'health', 'strength', 'health'],
    ['strength', 'speed', 'speed'],
    ['strength', 'speed', 'health', 'speed']
]

pts = [
    [1, 2, 1.5, -1],
    [2, -1.5, 1.5],
    [-1, 2, 0, 1.5]
]

s_cat = pd.Series(cat, name='cat')
s_pts = pd.Series(pts, name='pts')
df = pd.concat([s_cat, s_pts], axis=1)

# 将每行的列表拆分为多行，保留原始行索引
df_exploded = df.explode(['cat', 'pts'], ignore_index=False)
df_exploded['pts'] = df_exploded['pts'].astype(float)

# 计算每个原始行下各分类的分数总和
sum_stats = df_exploded.groupby([df_exploded.index, 'cat'])['pts'].sum().unstack()
sum_stats = sum_stats.rename(columns=lambda x: f"{x}_sum")

# 计算每个原始行下各分类的正分数计数
pos_counts = df_exploded[df_exploded['pts'] > 0].groupby([df_exploded.index, 'cat']).size().unstack(fill_value=0)
pos_counts = pos_counts.rename(columns=lambda x: f"{x}_pos")

# 计算每个原始行下各分类的负分数计数
neg_counts = df_exploded[df_exploded['pts'] < 0].groupby([df_exploded.index, 'cat']).size().unstack(fill_value=0)
neg_counts = neg_counts.rename(columns=lambda x: f"{x}_neg")

# 合并所有统计结果到原始DataFrame
result_df = pd.concat([df, sum_stats, pos_counts, neg_counts], axis=1)

# 给未出现的分类总和列填充NaN，匹配期望格式
for col in ['speed_sum', 'health_sum', 'strength_sum']:
    if col not in result_df.columns:
        result_df[col] = pd.NA
result_df[['speed_sum', 'health_sum', 'strength_sum']] = result_df[['speed_sum', 'health_sum', 'strength_sum']].fillna(0).where(result_df[['speed_sum', 'health_sum', 'strength_sum']].notna(), pd.NA)

# 调整列顺序与期望结果一致
desired_cols = ['cat', 'pts', 'speed_sum', 'health_sum', 'strength_sum', 
               'speed_pos', 'speed_neg', 'health_pos', 'health_neg', 
               'strength_pos', 'strength_neg']
result_df = result_df[desired_cols]

print(result_df)

运行结果 #

cat              pts  speed_sum  health_sum  strength_sum  speed_pos  speed_neg  health_pos  health_neg  strength_pos  strength_neg
0  [speed, health, strength, health]  [1, 2, 1.5, -1]        1.0         1.0           1.5          1          0           1           1             1             0
1           [strength, speed, speed]   [2, -1.5, 1.5]        0.0         NaN           2.0          1          1           0           0             1             0
2   [strength, speed, health, speed]  [-1, 2, 0, 1.5]        3.5         0.0          -1.0          2          0           0           0             0             1

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