Search code examples
pythonpandasdataframeapply

Add 2 new columns to existing dataframe using apply


I want to use the apply function that: - Takes 2 columns as inputs - Outputs two new columns based on a function.

An example is with this add_multiply function.

#function with 2 column inputs and 2 outputs
def add_multiply (a,b):
  return (a+b, a*b )

#example dataframe
df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})

#this doesn't work
df[['add', 'multiply']] = df.apply(lambda x: add_multiply(x['col1'], x['col2']), axis=1)

ideal result:

col1  col2  add  multiply
1     3     4    3
2     4     6    8


Solution

  • You can add result_type='expand' in the apply:

    ‘expand’ : list-like results will be turned into columns.

    df[['add', 'multiply']]=df.apply(lambda x: add_multiply(x['col1'], x['col2']),axis=1,
                                 result_type='expand')
    

    Or call a dataframe constructor:

    df[['add', 'multiply']]=pd.DataFrame(df.apply(lambda x: add_multiply(x['col1'], 
                                        x['col2']), axis=1).tolist())
    

       col1  col2  add  multiply
    0     1     3    4         3
    1     2     4    6         8