How do you groupby and aggregate using conditional statements in Pandas?

Viewed 194

Expanding on the question here, I'm wondering how to add aggregation to the following based on conditions:

Index    Name    Item            Quantity
0        John    Apple Red       10
1        John    Apple Green      5
2        John    Orange Cali     12
3        Jane    Apple Red       10
4        Jane    Apple Green      5
5        Jane    Orange Cali     18
6        Jane    Orange Spain     2
7        John    Banana           3
8        Jane    Coconut          5
9        John    Lime            10
... And so forth

What I need to do is getting this data converted into a dataframe like the following. Note: I am only interested in getting the total quantity of the apples and oranges both of them in separate columns, i.e. whatever other items appear in a certain group are not to be included in the aggregation done on column "Quantity" (but they are still to appear in the column "All items" as strings):

Index    Name    All Items                                          Apples Total  Oranges Total
0        John    Apple Red, Apple Green, Orange Cali, Banana, Lime  15             12
1        Jane    Apple Red, Apple Green, Orange Cali, Coconut       15             20

How would do I achieve that? Many thanks in advance!

4 Answers

Edits: fixed a bug.

To do this, before doing your groupby you can create your Total columns. These will contain a the number of apples and oranges in that row, depending whether that row's Item is apple or orange.

df['Apples Total'] = df.apply(lambda x: x.Quantity if ('Apple' in x.Item) else 0, axis=1)
df['Oranges Total'] = df.apply(lambda x: x.Quantity if ('Orange' in x.Item) else 0, axis=1)

When this is in place, groupby name and aggregate on each column. Sum on the total columns, and aggregate to list on the item column.

df.groupby('Name').agg({'Apples Total': 'sum',
                        'Oranges Total': 'sum',
                        'Item': lambda x: list(x)
                        })

You can use groupby and pivot_table after extracting Apple and Orange sub strings as below:

import re
s = df['Item'].str.extract("(Apple|Orange)",expand=False,flags=re.I)
# re.I used above is optional and is used for case insensitive matching

a = df.assign(Item_1=s).dropna(subset=['Item_1'])

out = (a.groupby("Name")['Item'].agg(",".join).to_frame().join(
        a.pivot_table("Quantity","Name","Item_1",aggfunc=sum).add_suffix("_Total"))
       .reset_index())

print(out)
   Name                                            Item  Apple_Total  \
0  Jane  Apple Red,Apple Green,Orange Cali,Orange Spain           15   
1  John               Apple Red,Apple Green,Orange Cali           15   

   Orange_Total  
0            20  
1            12  

EDIT: For edited question, you can use the same code only except groupby on the original dataframe df instead of the subset a and then join:

out = (df.groupby("Name")['Item'].agg(",".join).to_frame().join(
        a.pivot_table("Quantity","Name","Item_1",aggfunc=sum).add_suffix("_Total"))
       .reset_index())
print(out)

   Name                                               Item  Apple_Total  \
0  Jane  Apple Red,Apple Green,Orange Cali,Orange Spain...           15   
1  John      Apple Red,Apple Green,Orange Cali,Banana,Lime           15   

   Orange_Total  
0            20  
1            12  

First Filter only the required rows using str.contains on the column Item

from io import StringIO
import pandas as pd

s = StringIO("""Name;Item;Quantity
John;Apple Red;10
John;Apple Green;5
John;Orange Cali;12
Jane;Apple Red;10
Jane;Apple Green;5
Jane;Orange Cali;18
Jane;Orange Spain;2
John;Banana;3
Jane;Coconut;5
John;Lime;10
""")

df = pd.read_csv(s,sep=';')

req_items_idx = df[df.Item.str.contains('Apple|Orange')].index

df_filtered = df.loc[req_items_idx,:]

Once you have them you can further pivot the data to get the required values based on Name

pivot_df = pd.pivot_table(df_filtered,index=['Name'],columns=['Item'],aggfunc='sum')
pivot_df.columns = pivot_df.columns.droplevel()
pivot_df.columns.name = None
pivot_df = pivot_df.reset_index()

Generate the Totals for Apples and Oranges

orange_columns = pivot_df.columns[pivot_df.columns.str.contains('Orange')].tolist()
apple_columns = pivot_df.columns[pivot_df.columns.str.contains('Apple')].tolist()

pivot_df['Apples Total'] = pivot_df.loc[:,apple_columns].sum(axis=1)
pivot_df['Orange Total'] = pivot_df.loc[:,orange_columns].sum(axis=1)

A wrapper function to combine the Items together

def combine_items(inp,columns):
    res = []
    for val,col in zip(inp.values,columns):
        if not pd.isnull(val):
           res += [col]
    return ','.join(res)

req_columns = apple_columns+orange_columns
pivot_df['Items'] = pivot_df[apple_columns+orange_columns].apply(combine_items,args=([req_columns]),axis=1)

Finally you can get the required columns in a single place and print the values

total_columns = pivot_df.columns[pivot_df.columns.str.contains('Total')].tolist()
name_item_columns = pivot_df.columns[pivot_df.columns.str.contains('Name|Items')].tolist()

>>> pivot_df[name_item_columns+total_columns]
   Name                                           Items  Apples Total  Orange Total
0  Jane  Apple Green,Apple Red,Orange Cali,Orange Spain          15.0          20.0
1  John               Apple Green,Apple Red,Orange Cali          15.0          12.0

The answer is intended to outline the individual steps and approach one can take to solve something similar to this

df = pd.read_csv(StringIO("""
Index,Name,Item,Quantity
0,John,Apple Red,10
1,John,Apple Green,5
2,John,Orange Cali,12
3,Jane,Apple Red,10
4,Jane,Apple Green,5
5,Jane,Orange Cali,18
6,Jane,Orange Spain,2
7,John,Banana,3
8,Jane,Coconut,5
9,John,Lime,10
"""))

Getting list of items

grouping by name to get the list of items

items_list = pd.DataFrame(df.groupby(["Name"])["Item"].apply(list)).rename(columns={"Item": "All Items"})
items_list
        All Items
Name    
Jane    [Apple Red, Apple Green, Orange Cali, Orange Spain, Coconut]
John    [Apple Red, Apple Green, Orange Cali, Banana, Lime]

getting count of name item groups

renaming the temp df items column such that all the apples/oranges are treated similarly

temp2 = df.groupby(["Name", "Item"])['Quantity'].apply(sum)
temp2 = pd.DataFrame(temp2).reset_index().set_index("Name")
temp2['Item'] = temp2['Item'].str.replace(r'(?:.*)(apple|orange)(?:.*)', r'\1', case=False,regex=True)
    temp2
    Item    Quantity
Name        
Jane    Apple   5
Jane    Apple   10
Jane    Coconut 5
Jane    Orange  18
Jane    Orange  2
John    Apple   5
John    Apple   10
John    Banana  3
John    Lime    10
John    Orange  12

getting the required pivot table

pivot table for getting items count as separate column and retaining just apple orange count

pivot_df = pd.pivot_table(temp2, values='Quantity', columns='Item', index=["Name"], aggfunc=np.sum)
pivot_df = pivot_df[['Apple', 'Orange']]
    pivot_df
    Item    Apple   Orange
    Name        
    Jane    15.0    20.0
    John    15.0    12.0

merging the items list df and the pivot_df

output = items_list.merge(pivot_df, on="Name").rename(columns = {'Apple': 'Apples
Total', 'Orange': 'Oranges Total'})
output
All Items   Apples Total    Oranges Total
Name            
Jane    [Apple Red, Apple Green, Orange Cali, Orange Spain, Coconut]    15.0    20.0
John    [Apple Red, Apple Green, Orange Cali, Banana, Lime] 15.0    12.0
Related