I found this super confusing but eventually figured out a way of achieving this that didn't hurt my brain. Here it is, sorry if it doesn't match the example well...
dataframe with no index
# function to do the calcs
def f(row):
my_a = row['a'] # row is a Series, my_a is a scalar string
if my_a == 'a': # dummy logic to calc new values based on the row values
return [1, 2] # return 2 values to update 2 columns
else:
return [4, 5]
# simple test frame
input = pd.DataFrame.from_dict({
'a': ['a', 'd'],
'b': ['b', 'e'],
'c': ['c', 'f'],
'x': [0, 0],
'y': [0, 0]
})
# apply the function to update the x and y columns with the returned values
input[['x','y']] = input.apply(f, axis=1)
dataframe with an index
if your dataframe has an index.. you need to be a bit more explicit when you are doing the apply to ensure that "list-like results will be turned into columns"...
def f(row): # function to do the calcs
my_a = row['a'] # row is a Series, my_a is a scalar string
my_index = row.name # you might also want to use the index value in the calcs
if my_a == 'a': # dummy logic to calc new values based on the row values
return [1, 2] # return 2 values to update 2 columns
else:
return [4, 5]
input = pd.DataFrame.from_dict({
'an_index': ['indx1', 'indx2'],
'a': ['a', 'd'],
'b': ['b', 'e'],
'c': ['c', 'f'],
'x': [0, 0],
'y': [0, 0]
}).set_index(['an_index'])
# apply the function to update the x and y columns with the returned values
input[['x','y']] = input.apply(f, axis=1, result_type='expand')