CODEWITHSUNDEEP
python-3.xpandasdataframe
Salim Fadhley
Salim Fadhley

Reputation: 8195

Can I use pandas.DataFrame.apply() to make more than one column at once?

I have been given a function that takes the values in a row in a dataframe and returns a bunch of additional values.

It looks a bit like this:

my_func(row) -> (value1, value2, value3... valueN)

I'd like each of these values to become assigned to new columns in my dataframe. Can I use DataFrame.apply() do add multiple columns in one go, or do I have to add columns one at a time?

It's obvious how I can use apply to generate one column at a time:

from eunrg_utils.testing_functions.random_dataframe import random_dataframe
df = random_dataframe(columns=2,rows=4)
df["X"] = df.apply(axis=1, func=lambda row:(row.A + row.B))
df["Y"] = df.apply(axis=1, func=lambda row:(row.A - row.B))

But what if the two columns I am adding are something that are more easily calculated together? In this case, I already have a function that gives me everything I need in one shot. I'd rather not have to call it multiple times or add a load of caching.

Is there a syntax I can use that would allow me to use apply to generate 2 columns at the same time? Something like this, but less broken:

# Broken Pseudo-code
from eunrg_utils.testing_functions.random_dataframe import random_dataframe
df = random_dataframe(columns=2,rows=4)
df["X"], df["Y"] = df.apply(axis=1, func=lambda row:(row.A + row.B, row.B-row.A))

What is the correct way to do something like this?

Upvotes: 1

Views: 36

Answers (1)

jezrael
jezrael

Reputation: 862691

You can assign list of columns names like:

df = pd.DataFrame(np.random.randint(10, size=(2,2)),columns=list('AB'))
df[["X", "Y"]] = df.apply(axis=1, func=lambda row:(row.A + row.B, row.B-row.A))

print (df)
   A  B   X  Y
0  2  8  10  7
1  4  3   6 -1

Upvotes: 1

Related Questions