reuse function with multiple string values pandas

user_10964685 0
Viewed 39

I'm hoping to streamline a function that only return columns based on a single string value. Using below, I have two distinct colours in a df. I want to pass each colour to a function. But I only want the output to include columns relating to that colour.

If I have numerous colours and multiple outputs within the function, the returned df gets too large.

import pandas as pd
import numpy as np

d = ({
    'Date' : ['1/1/18','1/1/18','2/1/18','3/1/18','1/2/18','1/3/18','2/1/19','3/1/19'],                 
    'Val' : ['A','B','C','D','A','B','C','D'],   
    'Blue' : ['Blue', 'Blue', 'Blue', np.NaN, np.NaN, 'Blue', np.NaN, np.NaN],   
    'Red' : [np.NaN, np.NaN, np.NaN, 'Red', 'Red', np.NaN, 'Red', 'Red']                                 
    })

df = pd.DataFrame(data = d)

df['Date'] = pd.to_datetime(df['Date'], format = '%d/%m/%y')

df['Count'] = df.Date.map(df.groupby('Date').size())


def func(df, val):
    
    df['%s_cat' % val] = df['Count'] * 2

    return df


blue = func(df, 'Blue')


red = func(df, 'Red')

Intended output (Blue):

        Date Val  Blue   Count  Blue_cat 
0 2018-01-01   A  Blue       2         4        
1 2018-01-01   B  Blue       2         4        
2 2018-01-02   C  Blue       1         2               
5 2018-03-01   B  Blue       1         2

Intended output (Red):

        Date Val  Blue  Red  Count   Red_cat
3 2018-01-03   D   NaN  Red      1         2
4 2018-02-01   A   NaN  Red      1         2
6 2019-01-02   C   NaN  Red      1         2
7 2019-01-03   D   NaN  Red      1         2
pandas python
1 Answers

Use boolean indexing with DataFrame.copy for avoid SettingWithCopyWarning, because if you modify values in filtered DataFrame later you will find that the modifications do not propagate back to the original data, and that Pandas does warning:

def func(df, val):
    df = df[df[val].eq(val)].copy()
    df[f'{val}_cat'] = df['Count'] * 2

    return df

blue = func(df, 'Blue')
print (blue)
        Date Val  Blue  Red  Count  Blue_cat
0 2018-01-01   A  Blue  NaN      2         4
1 2018-01-01   B  Blue  NaN      2         4
2 2018-01-02   C  Blue  NaN      1         2
5 2018-03-01   B  Blue  NaN      1         2

red = func(df, 'Red')
print (red)
        Date Val Blue  Red  Count  Red_cat
3 2018-01-03   D  NaN  Red      1        2
4 2018-02-01   A  NaN  Red      1        2
6 2019-01-02   C  NaN  Red      1        2
7 2019-01-03   D  NaN  Red      1        2
user_29010020
Related