CODEWITHSUNDEEP
pythonpandas
Kumar AK
Kumar AK

Reputation: 1037

Selecting rows before and after rows based on value of other group--pandas

import pandas as pd
import numpy as np

raw_data = {'Country':['UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK','UK'],
    'Product':['A','A','A','A','B','B','B','B','B','B','B','B'],
            'Week': [1,2,3,4,1,2,3,4,5,6,7,8], 
       'val': [5,4,3,1,5,6,7,8,9,10,11,12]
    }

have = pd.DataFrame(raw_data, columns = ['Country','Product','Week', 'val'])

print(have)

enter image description here

i want to select rows before and after of last week of product A. i.e.week number 4 and product B rows should be 4 rows and before week was 3 and after weeks are 5,6 and including week 4 so total 4. Here's the wanted output

enter image description here

Upvotes: 0

Views: 45

Answers (1)

Henry Yik
Henry Yik

Reputation: 22503

IIUC:

df = pd.DataFrame(raw_data, columns=['Country', 'Product', 'Week', 'val'])

max_week = df.loc[df["Product"].eq("A"),"Week"].max()

print (df[df["Product"].eq("A")|((df["Week"]>=max_week-1)&(df["Week"]<=max_week+2))])

#
  Country Product  Week  val
0      UK       A     1    5
1      UK       A     2    4
2      UK       A     3    3
3      UK       A     4    1
6      UK       B     3    7
7      UK       B     4    8
8      UK       B     5    9
9      UK       B     6   10

Upvotes: 1

Related Questions