Groupby and subsetting rows based on condition

Viewed 286

I want to filter my dataframe.

My dataframe

  Col1    col2
0    A  event1
1    A  event2
2    A  event3
3    A  event2
4    B  event1
5    B  event3
6    B  event2
7    B  event2

The output dataframe

  Col1    col2
     A  event1
     B  event1
     B  event3

It should return the rows before event2 for each group. so far i tried

df.groupby('col1').apply(lambda x :x[0:x[x['col2'] == 'event2'].index[0]])

But it's not returning the required rows.

1 Answers

We can do this with groupby and cumsum, followed by a final filtering step:

df[df.col2.eq('event2').groupby(df.Col1).cumsum().eq(0)]

  Col1    col2
0    A  event1
4    B  event1
5    B  event3

To reset the index to a monotonically increasing range, use

df[df.col2.eq('event2').groupby(df.Col1).cumsum().eq(0)].reset_index(drop=True)

  Col1    col2
0    A  event1
1    B  event1
2    B  event3

Scott Boston suggests a nice improvement to the solution above using cumprod on a boolean mask. The principle is the same, but is more clean:

df[df.col2.ne('event2').groupby(df.Col1).cumprod()]

  Col1    col2
0    A  event1
4    B  event1
5    B  event3

groupby + idxmax based filtering suggested by W-B:

df[df.index < df.col2.eq('event2').groupby(df.Col1).transform('idxmax')]

  Col1    col2
0    A  event1
4    B  event1
5    B  event3
Related