python - 如果满足 groupby sum 条件，则 Pandas 追加

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]

但这不起作用。

标签： pythonpandas

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