Flagging NaN values based on a condition and year-CodePudding

I am trying to get this requirement of flagging NaN values based on condition and particular year, below is my code:

import pandas as pd
import numpy as np
s={'Fruits':['Apple','Orange', 'Banana', 'Mango'],'month':['201401','201502','201603','201604'],'weight':[2,4,1,6],'Quant':[251,178,298,300]}
p=pd.DataFrame(data=s)
upper = 250

How would I be able to flag NaN values for month- 201603 and 201604 (03 and 04 are the months), if upper>250. Basically my intention is to check if Quant value is greater than defined upper value, but for specific date i.e. 201603 and 201604.

This is how the output should look like-

Fruits  month   weight  Quant
0   Apple   201401  2   251.0
1   Orange  201502  4   178.0
2   Banana  201603  1   NaN
3   Mango   201604  6   NaN

CodePudding user response：

You could build a boolean condition that checks if "Quant" is greater than "upper" and the month is "03" or "04", and mask "Quant" column:

p['Quant'] = p['Quant'].mask(p['Quant'].gt(upper) & p['month'].str[-2:].isin(['03','04']))

Output:

   Fruits   month  weight  Quant
0   Apple  201401       2  251.0
1  Orange  201502       4  178.0
2  Banana  201603       1    NaN
3   Mango  201604       6    NaN

CodePudding user response：

Use:

p['Quant1'] = p[~(((p['month']=='201603')|(p['month']=='201604'))&(p['Quant']>250))]['Quant']

CodePudding user response：

You can use .loc:

p.loc[(p.Quant > upper) & (p.month.str[-2:].isin(['03','04'])), 'Quant'] = np.nan

OutPut:

   Fruits   month  weight  Quant
0   Apple  201401       2  251.0
1  Orange  201502       4  178.0
2  Banana  201603       1    NaN
3   Mango  201604       6    NaN