Filter dataframe by 5 rows above on condition

user_14114654 0
Viewed 54

I have a large df of purchases, sorted by person and datetime. I want to shorten the df to rows where a fruit is purchased and up to 5 rows before that fruit purchase for each person. If one of the 5 rows is a fruit purchase, then I only need to extract the rows above till that last fruit purchase.

fruits = ["apple", "banana", "pear"]
  date    pesron  item
10:00      amy    apple
10:01      amy    pear
10:02      amy    ipad       (6 rows above 10:08 banana purchase, so delete)
10:03      amy    headphones
10:04      amy    missing
10:05      amy    laptop
10:06      amy    unknown
10:07      amy    table
10:08      amy    banana   (one of the 5 rows above 10:10 banana purchase is also a fruit so stop here)
10:09      amy    unknown
10:10      amy    banana
...        ben

Expected Output

10:00      amy    apple
10:01      amy    pear
10:03      amy    headphones
10:04      amy    missing
10:05      amy    laptop
10:06      amy    unknown
10:07      amy    table
10:08      amy    banana
10:09      amy    unknown
10:10      amy    banana

I tried

m1= df.item.isin(fruits)
df.loc[m1, df.iloc[idx-5,:]]
pandas group-by python
1 Answers

You can create helper groups by swap order by indexing of rows with cumulative sum and then for 5 rows above with original row use GroupBy.tail:

print (df)
     date pesron        item
0   10:00    amy       apple
1   10:01    amy        pear
2   10:02    amy        ipad
3   10:03    amy  headphones
4   10:04    amy     missing
5   10:05    amy      laptop
6   10:06    amy     unknown
7   10:07    amy       table
8   10:08    amy      banana
9   10:09    amy     unknown
10  10:10    amy         ban <- change row
11  10:09    ben     unknown <- added row

fruits = ["apple", "banana", "pear"]

#swap rows and test for mmbership of list
m = df.iloc[::-1].item.isin(fruits)
#group by column and helper Series by cumulative sum of mask
g = m.groupby([df['pesron'], m.cumsum()])

#remove all rows after last match
mask1 = g.cummax()

#counter of rows per groups
mask2 = g.cumcount().lt(7).iloc[::-1]

#chain masks
df = df[mask1 & mask2]
print (df)
    date pesron        item
0  10:00    amy       apple
1  10:01    amy        pear
2  10:02    amy        ipad
3  10:03    amy  headphones
4  10:04    amy     missing
5  10:05    amy      laptop
6  10:06    amy     unknown
7  10:07    amy       table
8  10:08    amy      banana
user_29010020
Related