Idiomatic way to create pandas dataframe as concatenation of function of another's rows

user_959926 0
Viewed 72

Say I have one dataframe

import pandas as pd
input_df = pd.DataFrame(dict(a=[1, 2], b=[2, 3]))

Also I have a function f that maps each row to another dataframe. Here's an example of such a function. Note that in general the function could take any form so I'm not looking for answers that use agg to reimplement the f below.

def f(row):
    return pd.DataFrame(dict(x=[row['a'] * row['b'], row['a'] + row['b']],
                             y=[row['a']**2, row['b']**2]))

I want to create one dataframe that is the concatenation of the function applied to each of the first dataframe's rows. What is the idiomatic way to do this?

output_df = pd.concat([f(row) for _, row in input_df.iterrows()])

I thought I should be able to use apply or similar for this purpose but nothing seemed to work.

   x  y
0  2  1
1  3  4
0  6  4
1  5  9
pandas dataframe python
2 Answers

You can use DataFrame.agg to calucalate sum and prod and numpy.ndarray.reshape, df.pow(2)/np.sqaure for calculating sqaure.

out = pd.DataFrame({'x': df.agg(['prod', 'sum'],axis=1).to_numpy().reshape(-1),
                    'y': np.square(df).to_numpy().reshape(-1)})
out

   x  y
0  2  1
1  3  4
2  6  4
3  5  9
user_124164530

Yoy should avoid iterating rows (How to iterate over rows in a DataFrame in Pandas).

Instead try:

df = df.assign(product=df.a*df.b, sum=df.sum(axis=1), 
    asq=df.a**2, bsq=df.b**2)

Then:

df = [[[p, s], [asq, bsq]] for p, s, asq, bsq in df.to_numpy()]
user_66928980
Related