Home In Pandas, how to retrieve the rows which created each group, after aggregation and filtering?

Questions

In Pandas, how to retrieve the rows which created each group, after aggregation and filtering?

March 27, 2023

Let

import pandas as pd

df = pd.DataFrame(
    {
        'a': ['A', 'A', 'B', 'B', 'B', 'C'],
        'b': [True, True, True, False, False, True]
    }
)

print(df)

groups = df.groupby('a')  # "A", "B", "C"
agg_groups = groups.agg({'b':lambda x: all(x)}) # "A": True, "B": False, "C": True
agg_df = agg_groups.reset_index()
filtered_df = agg_df[agg_df["b"]]  # "A": True, "C": True

print(filtered_df)


# Now I want to get back the original df's rows, but only the remaining ones after group filtering

current output:

   a      b
0  A   True
1  A   True
2  B   True
3  B  False
4  B  False
5  C   True
   a     b
0  A  True
2  C  True

Required:

   a      b
0  A   True
1  A   True
2  B   True
3  B  False
4  B  False
5  C   True
   a     b
0  A  True
2  C  True
   a      b
0  A   True
1  A   True
5  C   True

>Solution :

Use GroupBy.transform for get all Trues to mask with same size like original DataFrame, so possible use boolean indexing:

df1 = df[df.groupby('a')['b'].transform('all')]

#alternative
#f = lambda x: x.all()
#df1 = df[df.groupby('a')['b'].transform(f)]
print (df1)
   a     b
0  A  True
1  A  True
5  C  True

If want filter in aggregation function output is boolean Series and filter match indices mapped by original column a:

ids = df.groupby('a')['b'].all()

df1 = df[df.a.isin(ids.index[ids])]
print (df1)
   a     b
0  A  True
1  A  True
5  C  True

Your solution is similar with filter column b:

groups = df.groupby('a')
agg_groups = groups.agg({'b':lambda x: all(x)})

df1 = df[df.a.isin(agg_groups.index[agg_groups['b']])]
print (df1)
   a     b
0  A  True
1  A  True
5  C  True