How to aggregate goupby and discard the rows after appearing a certain value?

Viewed 53

Say I have a given dataframe as below

input = pd.DataFrame({"id":[1,1,1,2,2,3,3,3,3,3], "values":["l", "m", "c", "l", "l", "l", "l", "c","c", "c"]})

and I wanted to remove the extra transactions after "c" appear for an id. say for id 3, the 1st 2 values are "l" and after that all transactions are value c so I only want the 1st c.

output = pd.DataFrame({"id":[1,1,1,2,2,3,3,3], "values": ["l", "m", "c", "l", "l", "l", "l", "c"]})

I tried to do drop_duplicates on a group by but it is not working as per my expectation:

input.groupby("id").drop_duplicates("values")

3 Answers

Create a boolean mask where values equals c, then use DataFrame.groupby to group this mask on id, then transform it using cumsum, finally use this mask to filter the dataframe:

# Here 'df' is your 'input' dataframe
mask = df['values'].eq('c').groupby(df['id']).cumsum().gt(1)
df1 = df[~mask]

Result:

print(df1)

   id values
0   1      l
1   1      m
2   1      c
3   2      l
4   2      l
5   3      l
6   3      l
7   3      c

If need remove only c rows after first c per groups:

Use DataFrame.duplicated with appended new column with compare c for compare values per groups (so tested duplicated by id and c), chaining by original mask m and last filtering by inverse mask by ~:

m = df['values'].eq('c')
df = df[~(df.assign(c = m).duplicated(['id','c']) & m)]
print (df)
   id values
0   1      l
1   1      m
2   1      c
3   2      l
4   2      l
5   3      l
6   3      l
7   3      c
 

Or if need remove all rows after first c per groups:

Use GroupBy.cumsum with boolean mask for remove values after first c by filtering with Series.le in boolean indexing per groups:

df = pd.DataFrame({"id":[1,1,1,2,2,3,3,3,3,3], 
                   "values":["l", "m", "c", "l", "l", "l", "l", "c","c", "c"]})

df = df[df['values'].eq('c').groupby(df['id']).cumsum().le(1)]
print (df)
   id values
0   1      l
1   1      m
2   1      c
3   2      l
4   2      l
5   3      l
6   3      l
7   3      c

You can create a dict of the index where the first occurence is by

In [24]: first_occurence = input.groupby('id').apply(lambda _df: (_df['values'] == 'c').idxmax() if np.any(_df['values'] == 'c') else None).to_dict()                                                      

In [25]: first_occurence                                                                                                                                                                                   
Out[25]: {1: 2.0, 2: nan, 3: 7.0}

Here, you need to return None if no values are found, otherwise you do leave out the last value if there is no 'c' for an id.

Then you can use DataFrame.truncate like this:

In [28]: input.groupby('id').apply(lambda _df: _df.truncate(after=first_occurence[_df['id'][0]])).droplevel(0)                                                                                             
Out[28]: 
   id values
0   1      l
1   1      m
2   1      c
3   2      l
4   2      l
5   3      l
6   3      l
7   3      c
Related