CODEWITHSUNDEEP
pythonpandasdataframegroup-by
Owad A.
Owad A.

Reputation: 1

How to create conditional group-by with pandas?

Suppose I have a dataframe like this:

data = [['A', 'HIGH', 120, 200],
        ['A', 'MID', 350, 200],
        ['B', 'HIGH', 130, 100],
        ['B', 'HIGH', 70, 100],
        ['A', 'MID', 130, 200]]
df = pd.DataFrame(data, columns=['Category', 'Range', 'Total', 'Avg'])

Now, I want to create a Group By that when the category is A, it groups by category and Range, while when it is B, it group only by category.

Is it possible to do?

Thanks!

Upvotes: 0

Views: 267

Answers (1)

Abhishek
Abhishek

Reputation: 1625

Check below code. It will also work B has multiple range.

import pandas as pd
import numpy as np 

data = [['A', 'HIGH', 120, 200],
        ['A', 'MID', 350, 200],
        ['A', 'MID', 130, 200],
        ['B', 'HIGH', 130, 100],
        ['B', 'MID', 70, 100],
        ['B', 'MID', 70, 100]
        ]
df = pd.DataFrame(data, columns=['Category', 'Range', 'Total', 'Avg'])

df[['Total_New','Avg_New']] = df.assign(group_col = np.where(df['Category']=='A',df.Category+df.Range, df.Category)).\
groupby('group_col')['Total','Avg'].transform('sum')

df

Output:

enter image description here

Upvotes: 1

Related Questions