如何按用户计算合并运费均值且仅当存在shipping_cost时生效

原始DataFrame #

index  user  default_shipping_cost     category  shipping_cost  shipping_coalesce  estimated_shipping_cost
0      0     1                      1      clothes            NaN                1.0                      6.0
1      1     1                      1  electronics            2.0                2.0                      6.0
2      2     1                     15    furniture            NaN               15.0                      6.0
3      3     2                     15    furniture            NaN               15.0                     15.0
4      4     2                     15    furniture            NaN               15.0                     15.0

需求说明 #

按用户将shipping_cost与default_shipping_cost合并，计算合并后运费的均值，但仅当该用户至少有一个非空的shipping_cost值时才执行计算；若用户没有任何非空的shipping_cost值，则对应结果为NaN。

• 用户1存在至少一个非空的shipping_cost值，因此计算均值
• 用户2没有任何非空的shipping_cost值，因此结果为NaN

现有代码（存在问题） #

import pandas as pd

pd.set_option("display.max_columns", None)
pd.set_option("display.max_rows", None)
pd.set_option('display.width', 1000)

df = pd.DataFrame(
    {
        'user': [1, 1, 1, 2, 2],
        'default_shipping_cost': [1, 1, 15, 15, 15],
        'category': ['clothes', 'electronics', 'furniture', 'furniture', 'furniture'],
        'shipping_cost': [None, 2, None, None, None]
    }
)
df.reset_index(inplace=True)
df['shipping_coalesce'] = df.shipping_cost.combine_first(df.default_shipping_cost)

dfg_user = df.groupby(['user'])
df['estimated_shipping_cost'] = dfg_user['shipping_coalesce'].transform("mean")
print(df)

上述代码会给用户2也计算出均值15.0，不符合需求。

修正后的代码 #

方法一：使用transform直接处理 #

import pandas as pd

pd.set_option("display.max_columns", None)
pd.set_option("display.max_rows", None)
pd.set_option('display.width', 1000)

df = pd.DataFrame(
    {
        'user': [1, 1, 1, 2, 2],
        'default_shipping_cost': [1, 1, 15, 15, 15],
        'category': ['clothes', 'electronics', 'furniture', 'furniture', 'furniture'],
        'shipping_cost': [None, 2, None, None, None]
    }
)
df.reset_index(inplace=True)
df['shipping_coalesce'] = df.shipping_cost.combine_first(df.default_shipping_cost)

# 标记每个用户是否有非空的shipping_cost
df['has_shipping_cost'] = df.groupby('user')['shipping_cost'].transform(lambda x: x.notna().any())

# 根据标记计算均值，否则设为NaN
df['estimated_shipping_cost'] = df.groupby('user')['shipping_coalesce'].transform(
    lambda x: x.mean() if df.loc[x.index, 'has_shipping_cost'].iloc[0] else pd.NA
)

# 删除辅助列
df.drop('has_shipping_cost', axis=1, inplace=True)

print(df)

方法二：先聚合用户统计再合并 #

import pandas as pd

pd.set_option("display.max_columns", None)
pd.set_option("display.max_rows", None)
pd.set_option('display.width', 1000)

df = pd.DataFrame(
    {
        'user': [1, 1, 1, 2, 2],
        'default_shipping_cost': [1, 1, 15, 15, 15],
        'category': ['clothes', 'electronics', 'furniture', 'furniture', 'furniture'],
        'shipping_cost': [None, 2, None, None, None]
    }
)
df.reset_index(inplace=True)
df['shipping_coalesce'] = df.shipping_cost.combine_first(df.default_shipping_cost)

# 聚合用户的两个关键信息：是否有shipping_cost、合并运费的均值
user_stats = df.groupby('user').agg(
    has_shipping_cost=('shipping_cost', lambda x: x.notna().any()),
    mean_coalesce=('shipping_coalesce', 'mean')
)

# 根据是否有shipping_cost决定最终均值，否则设为NaN
user_stats['estimated_shipping_cost'] = user_stats.apply(
    lambda row: row['mean_coalesce'] if row['has_shipping_cost'] else pd.NA, axis=1
)

# 将结果合并回原DataFrame
df = df.merge(user_stats[['estimated_shipping_cost']], on='user', how='left')

print(df)

预期输出 #

index  user  default_shipping_cost     category  shipping_cost  shipping_coalesce estimated_shipping_cost
0      0     1                      1      clothes            NaN                1.0                      6.0
1      1     1                      1  electronics            2.0                2.0                      6.0
2      2     1                     15    furniture            NaN               15.0                      6.0
3      3     2                     15    furniture            NaN               15.0                     <NA>
4      4     2                     15    furniture            NaN               15.0                     <NA>

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