Python pandas - filter rows after groupby

Question

For example I have following table:

index,A,B
0,0,0
1,0,8
2,0,8
3,1,0
4,1,5

After grouping by A:

0:
index,A,B
0,0,0
1,0,8
2,0,8

1:
index,A,B
3,1,5
4,1,3

What I need is to drop rows from each group, where the number in column B is less than maximum value from all rows from group's column B. Well I have a problem translating and formulating this problem to English so here is the example:

Maximum value from rows in column B in group 0: 8

So I want to drop row with index 0 and keep rows with indexes 1 and 2

Maximum value from rows in column B in group 1: 5

So I want to drop row with index 4 and keep row with index 3

I have tried to use pandas filter function, but the problem is that it is operating on all rows in group at one time:

data = <example table>
grouped = data.groupby("A")
filtered = grouped.filter(lambda x: x["B"] == x["B"].max())

So what I ideally need is some filter, which iterates through all rows in group.

Thanks for help!

P.S. Is there also way to only delete rows in groups and do not return DataFrame object?

Answer 1

You just need to use apply on the groupby object. I modified your example data to make this a little more clear:

import pandas
from io import StringIO

csv = StringIO("""index,A,B
0,1,0.0
1,1,3.0
2,1,6.0
3,2,0.0
4,2,5.0
5,2,7.0""")

df = pandas.read_csv(csv, index_col='index')
groups = df.groupby(by=['A'])
print(groups.apply(lambda g: g[g['B'] == g['B'].max()]))

Which prints:

         A  B
A index      
1 2      1  6
2 4      2  7