如何在pd.groupby中使用np.where为分组分配新变量？

示例数据 #

d = {'id' : ["A","A","A","A","A","A","B","B","B","B","B","B"],
    'month' : [1,1,1,1,2,2,1,1,1,1,2,2],
    'week' : [1,2,3,4,1,2,1,2,3,4,1,2]}

example_df = pd.DataFrame(data = d)

尝试的代码及错误 #

尝试的代码：

example_df = (example_df.assign(has_3_or_4 = example_df.groupby(['id', 'month'])
                .apply(lambda x: np.where(any(x.week.isin([3, 4])),'has_3_or_4', 'no_3_or_4'))))

返回错误：

TypeError: incompatible index of inserted column with frame index

添加as_index=False后代码可运行，但结果不符合预期。

R中的实现方式 #

example_df %>% group_by(id, month) %>% mutate(has_3_or_4 = if_else(any(c(3,4) %in% week), 'has_3_or_4', 'no_3_or_4'))

解决方案 #

方法1：用transform替代apply（最贴合R逻辑） #

apply返回的是以分组键为索引的结果，和原DataFrame的行索引不匹配，所以报错。transform会把每组的结果广播到组内每一行，完美匹配原数据结构：

import numpy as np
import pandas as pd

def check_week(group):
    return 'has_3_or_4' if any(group.isin([3,4])) else 'no_3_or_4'

example_df['has_3_or_4'] = example_df.groupby(['id', 'month'])['week'].transform(check_week)

方法2：先计算分组结果再合并 #

先单独算出每个分组的标记值，再通过关联合并到原数据：

# 计算每个(id, month)分组的标记结果
group_tags = example_df.groupby(['id', 'month'])['week']\
    .apply(lambda x: 'has_3_or_4' if x.isin([3,4]).any() else 'no_3_or_4')\
    .rename('has_3_or_4')

# 合并到原DataFrame
example_df = example_df.merge(group_tags, on=['id', 'month'], how='left')

方法3：用groupby.any()生成映射 #

先判断每组是否包含目标值，再通过索引映射赋值：

# 生成分组标记的布尔值，再转成目标字符串
group_bool = example_df.groupby(['id', 'month'])['week'].apply(lambda x: x.isin([3,4]).any())
tag_map = group_bool.map({True: 'has_3_or_4', False: 'no_3_or_4'})

# 给原数据赋值
example_df['has_3_or_4'] = example_df.set_index(['id', 'month']).index.map(tag_map)

以上三种方法都能实现和R中mutate完全一致的效果，其中方法1的逻辑最直观，和R的写法思路最接近。

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