rgh_dsa
rgh_dsa

Reputation: 117

New column based off certain input parameter to select what columns to use - Python

Have a pandas dataframe that includes multiple columns of monthly finance data. I have an input of period that is specified by the person running the program. It's currently just saved as period like shown below within the code.

#coded into python
period = ?? (user adds this in from input screen)

I need to create another column of data that uses the input period number to perform a calculation of other columns.

enter image description here

So, in the above table I'd like to create a new column 'calculation' that depends on the period input. For example, if a period of 1 was used the following calc1 would be completed (with math actually done). Period = 2 - then calc2. Period = 3 - then calc3. I only need one column calculated depending on the period number but added three examples in below picture for example of how it'd work.

enter image description here

I can do this in SQL using case when. So using the input period then sum what columns I need to.

select  Account #,
'&Period' AS Period,                    
'&Year' AS YR,                  
case                    
  When '&Period' = '1' then sum(d_cf+d_1)                   
  when '&Period' = '2' then sum(d_cf+d_1+d_2)                   
  when '&Period' = '3' then sum(d_cf+d_1+d_2+d_3)                   

I am unsure on how to do this easily in python (newer learner). Yes, I could create a column that does each calculation via new column for every possible period (1-12), and then only select that column but I'd like to learn and do it a more efficient way.

Can you help more or point me in a better direction?

Upvotes: 0

Views: 361

Answers (3)

Mark Moretto
Mark Moretto

Reputation: 2348

Setup:

import pandas as pd

ddict = {
    'Year':['2018','2018','2018','2018','2018',],
    'Account_Num':['1111','1122','1133','1144','1155'],
    'd_cf':['1','2','3','4','5'],
    }

data = pd.DataFrame(ddict)

Create value calculator:

def get_calcs(period):
    # Convert period to integer
    s = str(period)

    # Convert to string value
    n = int(period) + 1

    # This will repeat the period number by the value of the period number
    return ''.join([i * n for i in s])

Main function copies data frame, iterates through period values, and sets calculated values to the correct spot index-wise for each relevant column:

def process_data(data_frame=data, period_column='d_cf'):
    # Copy data_frame argument
    df = data_frame.copy(deep=True)

    # Run through each value in our period column
    for i in df[period_column].values.tolist():

        # Create a temporary column
        new_column = 'd_{}'.format(i)

        # Pass the period into our calculator; Capture the result
        calculated_value = get_calcs(i)

        # Create a new column based on our period number
        df[new_column] = ''

        # Use indexing to place the calculated value into our desired location
        df.loc[df[period_column] == i, new_column] = calculated_value

    # Return the result
    return df

Start:

   Year Account_Num d_cf
0  2018        1111    1
1  2018        1122    2
2  2018        1133    3
3  2018        1144    4
4  2018        1155    5

Result:

process_data(data)

   Year Account_Num d_cf d_1  d_2   d_3    d_4     d_5
0  2018        1111    1  11                          
1  2018        1122    2      222                     
2  2018        1133    3           3333               
3  2018        1144    4                 44444        
4  2018        1155    5                        555555

Upvotes: 1

N.Clarke
N.Clarke

Reputation: 268

You could certainly do something like

df[['d_cf'] + [f'd_{i}' for i in range(1, period+1)]].sum(axis=1)

Upvotes: 1

YOLO
YOLO

Reputation: 21749

You can do this using a simple function in python:

def get_calculation(df, period=NULL): 

    '''
    df = pandas data frame
    period = integer type
    '''

    if period == 1:
        return df.apply(lambda x: x['d_0'] +x['d_1'], axis=1)

    if period == 2:
        return df.apply(lambda x: x['d_0'] +x['d_1']+ x['d_2'], axis=1)

    if period == 3:
        return df.apply(lambda x: x['d_0'] +x['d_1']+ x['d_2'] + x['d_3'], axis=1)

new_df = get_calculation(df, period = 1)

Setup:

df = pd.DataFrame({'d_0':list(range(1,7)),
                   'd_1': list(range(10,70,10)),
                   'd_2':list(range(100,700,100)),
                   'd_3': list(range(1000,7000,1000))})

Upvotes: 1

Related Questions