如何在Pandas的transform中设置条件，按组统计并生成新列

原始数据 #

创建DataFrame的代码：

import pandas as pd
df = pd.DataFrame({
    "A": ["a1", "a1", "a1", "a1", "a1", "a2", "a2", "a2", "a3"],
    "B": ["a1", "a4", "a5", "a6", "a10", "a7", "a8", "a9", "a3"],
    "C": ["Hello", "World", "How", "are", "you", "today", "miss", "?", "!"]
})

原始表格：

A    B      C
0  a1   a1  Hello
1  a1   a4  World
2  a1   a5    How
3  a1   a6    are
4  a1  a10    you
5  a2   a7  today
6  a2   a8   miss
7  a2   a9      ?
8  a3   a3      !

需求说明 #

生成新列n，其中n为按A分组后，每组中满足A≠B的行数，且每组所有行都填充该计数（包括A=B的行），期望结果：

A    B      C   n
0  a1   a1  Hello   4
1  a1   a4  World   4
2  a1   a5    How   4
3  a1   a6    are   4
4  a1  a10    you   4
5  a2   a7  today   3
6  a2   a8   miss   3
7  a2   a9      ?   3
8  a3   a3      !   0

尝试的代码及问题 #

尝试以下代码：

df["n"] = df.loc[df.A != df.B].groupby("A")["B"].transform(len)

但结果中A=B的行出现NaN：

A    B      C     n
0  a1   a1  Hello   NaN
1  a1   a4  World   4.0
2  a1   a5    How   4.0
3  a1   a6    are   4.0
4  a1  a10    you   4.0
5  a2   a7  today   3.0
6  a2   a8   miss   3.0
7  a2   a9      ?   3.0
8  a3   a3      !   NaN

问题原因：先过滤了A≠B的行，再分组计算，原DataFrame中被过滤的行（A=B）无法匹配到分组结果，导致NaN。

解决方案 #

无需过滤原DataFrame，直接按A分组后计算每组满足A≠B的行数，再通过map将结果映射到每一行：

# 先计算每个A分组对应的满足条件的行数
group_counts = df.groupby("A").apply(lambda grp: (grp["A"] != grp["B"]).sum())
# 将计数映射到原DataFrame的每一行
df["n"] = df["A"].map(group_counts)

代码解释 #

1. groupby("A").apply(lambda grp: (grp["A"] != grp["B"]).sum())：对每个A分组，计算组内A≠B的行数（布尔值True等价于1，sum直接得到计数）。
2. df["A"].map(group_counts)：将每个行的A值对应的分组计数填充到n列，所有行都会得到有效值，不会出现NaN。

执行后即可得到期望结果。

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