CODEWITHSUNDEEP
pythonpandas
Leon Adams
Leon Adams

Reputation: 491

Create new pandas dataframe columns from function that returns tuple

I have a function that returns tuple:

def pwrs(x):
   return x*x, x*x*x, x*x*x*x

I would like to apply this function to a single column dataframe named data:

+-------+
|    x  |
+-------+
|    1  |
|    2  |
|    3  |
|    4  |
+-------+

and get back a new dataframe with new columns based on the function return tuples:

+---+------+------+------+
| x |  x^2 |  x^3 |  x^4 |
+---+------+------+------+
| 1 |    1 |    1 |    1 |
| 2 |    4 |    8 |   16 |
| 3 |    9 |   27 |   81 |
| 4 |   16 |   64 |  256 |
+---+------+------+------+

Got as far as iterating through the rows and applying the function:

for _, row in data.iterrows():
    print(pwrs(row['x']))

Not sure how to progress from here....

Upvotes: 3

Views: 447

Answers (4)

Dani Mesejo
Dani Mesejo

Reputation: 61910

You could do:

df[['x^2', 'x^3', 'x^4']] = df.x.apply(lambda x: pd.Series(pwrs(x)))
print(df)

Output

   x  x^2  x^3  x^4
0  1    1    1    1
1  2    4    8   16
2  3    9   27   81
3  4   16   64  256

Notice that this works for any function that returns a tuple, not just this mathematical operations.

Upvotes: 1

piRSquared
piRSquared

Reputation: 294258

Generalized

c = np.arange(1, 5)
pd.DataFrame(df.to_numpy() ** c, df.index, c).add_prefix('x^')

   x^1  x^2  x^3  x^4
0    1    1    1    1
1    2    4    8   16
2    3    9   27   81
3    4   16   64  256

Upvotes: 4

sacuL
sacuL

Reputation: 51335

If you want to use your function, you can simply use the dataframe column as the argument:

df['x^2'], df['x^3'], df['x^4'] = pwrs(df['x'])


   x  x^2  x^3  x^4
0  1    1    1    1
1  2    4    8   16
2  3    9   27   81
3  4   16   64  256

Upvotes: 2

BENY
BENY

Reputation: 323226

IIUC

pd.DataFrame(df.x.values[:,None]**np.array([1,2,3,4]))
Out[290]: 
   0   1   2    3
0  1   1   1    1
1  2   4   8   16
2  3   9  27   81
3  4  16  64  256

Upvotes: 3

Related Questions