Pandas apply on multiple columns and return multiple columns

Question

Here is an example to simplify the question:

import pandas as pd

data = {
  "X": [420, 380, 390],
  "Y": [50, 40, 45]
}
df = pd.DataFrame(data)

print(df) 

Let's say we want to make a simple operation:

if X > Y: new_X = X - 10  
if Y > = X: new_Y = Y - 10

def fct(x, y):
   if x>y:
    x = x -10
   else:
    y = y - 10
    return x, y
s = df.apply(lambda x: fct(x.X, x.Y), axis=1)
ss = s.explode().values.reshape(-1, 2)
df[['col1', 'col2']] = ss
        X   Y   col1 col2
    0  420  50  410  50 
    1  380  40  370  40 
    2  390  45  380  45
 

Is there any other optimized way? I tried with some more complicated data type such as dict or nested list and explode can fail

Answer 1

Based on your condition, try this code:

df.loc[df["X"] > df["Y"],  "col1"] = df["X"]-10
df.loc[df["X"] <= df["Y"],  "col1"] = df["X"]
df.loc[df["X"] <= df["Y"],  "col2"] = df["Y"]-10
df.loc[df["Y"] < df["X"],  "col2"] = df["Y"]