Returning multiple variables with pandas.series.apply-CodePudding

I want to apply a single function to a dataframe column. This function returns multiple results which I want to go to multiple columns in the original dataframe. I cant seem to get around the "too many values to unpack" error....

df = pd.DataFrame(data={'x': [1,2,3,4]})

   x
0  1
1  2
2  3
3  4

The function I want to apply :

def do_math(x):
    double = 2*x
    triple = 3*x
    return double, triple

My attempt to apply it (which doesnt work):

df['DOUBLE'], df['TRIPLE'] = df['x'].apply(do_math)

What I want :

   x  DOUBLE  TRIPLE
0  1       2       3
1  2       4       6
2  3       6       9
3  4       8      12

CodePudding user response：

You can convert the result of Series.apply to list then assign to multiple columns

df[['DOUBLE', 'TRIPLE']] = df['x'].apply(do_math).tolist()

print(df)

   x  DOUBLE  TRIPLE
0  1       2       3
1  2       4       6
2  3       6       9
3  4       8      12

You can also try DataFrame.apply on rows with result_type='expand'

df[['DOUBLE', 'TRIPLE']] = df.apply(lambda row: do_math(row['x']), axis=1, result_type='expand')

print(df)

   x  DOUBLE  TRIPLE
0  1       2       3
1  2       4       6
2  3       6       9
3  4       8      12

Since your operation is simple, you can also try df.eval

df = df.eval('''
double = 2 * x
triple = 3 * x
'''
)

print(df)

   x  double  triple
0  1       2       3
1  2       4       6
2  3       6       9
3  4       8      12

CodePudding user response：

I would recommend doing them separately:

df = pd.DataFrame(data={'x': [1,2,3,4]})
df['DOUBLE'] = df['x'].apply(lambda x: 2*x)
df['TRIPLE'] = df['x'].apply(lambda x: 3*x)

CodePudding user response：

You could change apply to return a Series:

def do_math(x):
    double = 2*x
    triple = 3*x
    return pd.Series([double, triple])

df[['DOUBLE', 'TRIPLE']] = df.apply(lambda row: do_math(row['x']), axis=1)

   x  DOUBLE  TRIPLE
0  1       2       3
1  2       4       6
2  3       6       9
3  4       8      12

CodePudding user response：

You can use df.transform when all the functions you want are defined separately.

def double(x):
    return x * 2

def triple(x):
    return x * 3

df[['DOUBLE', 'TRIPLE']] = df['x'].transform([double, triple])

This can take any number of functions you want. You can add more functions when needed and won't have to modify other functions.