如果满足groupby sum条件，熊猫会追加

Question

很抱歉，如果这是一个重复的问题，我不确定我要做什么的特定语法。

我想遍历大df，其中A和B是索引值，x，y，z是数据列

df=
A  B     x  y  z
0.1 0.2  2  2  0
0.1 0.3  1  3  0
0.1 0.4  3  3  0
0.2 0.2  4  1  -1
0.2 0.3  5  3  0
0.2 0.1  6  1  0
0.3 0.2  1  1  0
0.3 0.5  1  2  0
0.3 0.7  2  1  0

如果满足以下条件：

df.groupby('A')['z'].sum()==0

将整个groupby对象附加到新的df或生成满足此条件的所有groupby obj的df。

预期输出：

new_df=

A   B    x  y   z 
0.1 0.2  2  2   0
0.1 0.3  1  3   0
0.1 0.4  3  3   0
0.3 0.2  1  1   0
0.3 0.5  1  2   0
0.3 0.7  2  1   0

我正在尝试类似

new_df = df.loc[df.groupby('A')['z'].sum())==0]

但这不起作用。

Answer 1

结合使用groupby和level=0来进行transfrom，然后根据布尔值mask来完成数据帧的子集

mask = df.groupby(level=0)['z'].transform(lambda s: s.sum()==0)

new_df = df[mask].copy()

new_df
#          x  y  z
# A   B           
# 0.1 0.2  2  2  0
#     0.3  1  3  0
#     0.4  3  3  0
# 0.3 0.2  1  1  0
#     0.5  1  2  0
#     0.7  2  1  0