In pandas transform lambda function use if else to get 2nd largest value if largest is infinity

Viewed 256

I have a dataframe as below:

      country_code  count_date  confirmed_cases  growth_rate
0              AFG  2020-09-18          38872.0    94.019934
1              AFG  2020-09-17          38855.0    83.641975
2              AFG  2020-09-16          38815.0   108.154506
3              AFG  2020-09-15          38772.0   196.296296
4              AFG  2020-09-14          38716.0   196.894410
...            ...         ...              ...          ...
20451          ZWE  2020-06-03            206.0          NaN
20452          ZWE  2020-06-02            203.0          NaN
20453          ZWE  2020-06-01            178.0          NaN
20454          ZWE  2020-05-31            174.0          NaN
20455          ZWE  2020-05-30            149.0          NaN

I am trying to do this: df.groupby('count_date').growth_rate.transform(lambda x: x/x.max())

But some of the x.max() values are infinity and in such cases I want to use the 2nd largest value. How do I convert the below pseudo-code to actual code:

df.groupby('count_date').growth_rate.transform(lambda x: x/x.max() if x.max() != inf else x/<2nd largest>)

1 Answers

You can use a mask to drop inf values from x :

df.groupby('count_date').growth_rate.transform(lambda x: x/x[~np.isinf(x)].max())

Here is an example:

# Sample DataFrame
df = pd.DataFrame({
    'count_date': ['2020-09-17', '2020-09-17', '2020-09-17', '2020-09-17'],
    'growth_rate': [ 1, 3, 2, np.inf]
})

df.groupby('count_date').growth_rate.transform(lambda x: x/x[~np.isinf(x)].max())
0    0.333333
1    1.000000
2    0.666667
3         inf

As you can see, the scaling was performed using the largest non infinity value found in x.

Related