Pekka
Pekka

Reputation: 2448

Define two columns with one map in Pandas DataFrame

I have a function which returns a list of length 2. I would like to apply this function to one column in my dataframe and assign the result to two columns.

This actually works:

from pandas import *

def twonumbers(x):
    return [2*x, 3*x]

df = DataFrame([1,4,11],columns=['v1'])

concat([df,DataFrame(df['v1'].map(twonumbers).tolist(), columns=['v2','v3'])],axis=1)

But I am looking for a simpler way to do the last line above. Something like this:

df['v3'], df['v2'] = df['v1'].map(twonumbers)

Upvotes: 3

Views: 2336

Answers (1)

unutbu
unutbu

Reputation: 880339

import pandas as pd

def twonumbers(x):
    return [2*x, 3*x]

df = pd.DataFrame([1,4,11], columns=['v1'])
df['v2'], df['v3'] = twonumbers(df['v1'])

makes df look like this

   v1  v2  v3
0   1   2   3
1   4   8  12
2  11  22  33

Note: This relies on twonumbers being able to accept a Pandas Series as input, and returning a list of two Series as output.

Upvotes: 3

Related Questions