CODEWITHSUNDEEP
pythonpython-3.xpandas
blastoise
blastoise

Reputation: 119

Apply function in pandas with multiple args and multiple outputs

I have a function foo() which takes two arguments as two columns of pandas named year and month and return a list of four numbers.

  df['A'],df['B'],df['C'],df['D']=  df.apply(lambda x: foo(x.year, x.month), axis=1,result_type="expand")

It just gives me four columns named A,B,C,D with 0,1,2,3 respectively filled in them. What am I doing wrong?

Kindly don't answer for single variable outputs or single variable arguments. There are plenty of examples out there for that. Thankyou so much for the help.

Upvotes: 1

Views: 1074

Answers (1)

Shubham Sharma
Shubham Sharma

Reputation: 71689

You are incorrectly assigning the results of df.apply with optional parameter result_type='expand', instead, Use:

df[['A', 'B', 'C', 'D']] =  df.apply(lambda x: foo(x.year, x.month), axis=1, result_type="expand")

Consider the example demonstrating this,

df = pd.DataFrame({'col1': np.arange(5), 'col2': np.arange(5) * 2})

#print(df)
   col1  col2
0     0     0
1     1     2
2     2     4
3     3     6
4     4     8

Reproducing the problem situation,

df['A'], df['B'] = df.apply(lambda s: (s['col1']**2, s['col2']**2), axis=1, result_type='expand')

#print(df)
   col1  col2  A  B
0     0     0  0  1
1     1     2  0  1
2     2     4  0  1
3     3     6  0  1
4     4     8  0  1

The solution is,

df[['A', 'B']] = df.apply(lambda s: (s['col1']**2, s['col2']**2), axis=1, result_type='expand')

#print(df)
   col1  col2   A   B
0     0     0   0   0
1     1     2   1   4
2     2     4   4  16
3     3     6   9  36
4     4     8  16  64

OR:

df['A'], df['B'] = df.apply(
lambda s: (s['col1']**2, s['col2']**2), axis=1, result_type='expand').T.to_numpy()

#print(df)
   col1  col2   A   B
0     0     0   0   0
1     1     2   1   4
2     2     4   4  16
3     3     6   9  36
4     4     8  16  64

Upvotes: 1

Related Questions