1

I want to apply a function that returns multiple values to my data frame such that these values are collected in different columns. How do we achieve this?

Minimum Reproducible code:

import pandas as pd
df=pd.DataFrame({'col1':[1,2,3,4,5], 'col2':['a','b','c','d','e']})

Which returns the table:

col1 col2
0 1 a
1 2 b
2 3 c
3 4 d
4 5 e

Now I want one of the return values to be a list, so following the answer from this StackOverFlow question, I made an empty column astype(object)

import numpy as np
df['e']=np.nan
df['e']=df['e'].astype('object')

Following the accepted answer in this question I am returning the value like this:

def funct(a,b):
  c=str(a)+b
  d=b+str(a)
  return[c,d,[c,d]]

If I want to save the output into three columns c,d,e, I am trying:

df[['c','d','e']]=df.apply(lambda x: funct(x['col1'],x['col2']), axis=1)

Which gives the error:

ValueError: Must have equal len keys and value when setting with an iterable

If I run

df['c','d','e']=df.apply(lambda x: funct(x['col1'],x['col2']), axis=1)

The return value is:

| i|col1| col2 |(c, d, e)         |
|:-|:--:|------|------------------:
|0 |  1 |  a   |[1a, a1, [1a, a1]]|
|1 |  2 |  b   |[2b, b2, [2b, b2]]|
|2 |  3 |  c   |[3c, c3, [3c, c3]]|
|3 |  4 |  d   |[4d, d4, [4d, d4]]|
|4 |  5 |  e   |[5e, e5, [5e, e5]]|

How do I get:

|  | col1 |col2| c | d | e      |
|:-|:----:|:--:|:-:|:-:|------- |
|0 | 1    |a   |1a |a1 |[1a, a1]|
|1 | 2    |b   |2b |b2 |[2b, b2]|
|2 | 3    |c   |3c |c3 |[3c, c3]|
|3 | 4    |d   |4d |d4 |[4d, d4]|
|4 | 5    |e   |5e |e5 |[5e, e5]|
Anirban Saha
  • 1,350
  • 2
  • 10
  • 38

1 Answers1

2

You can do result_type='expand' in apply with your existing function:

df[['c','d','e']]=(df.apply(lambda x: funct(x['col1'],x['col2']),
                  axis=1,result_type='expand')

print(df)

   col1 col2   c   d         e
0     1    a  1a  a1  [1a, a1]
1     2    b  2b  b2  [2b, b2]
2     3    c  3c  c3  [3c, c3]
3     4    d  4d  d4  [4d, d4]
4     5    e  5e  e5  [5e, e5]
anky
  • 74,114
  • 11
  • 41
  • 70