4

Say I have a function:

def fn(x)
    y = x ** 2
    z = x ** 3
    return y, z

And I want to use df['x'].apply(lambda x: fn(x)) to return both y and z in separate columns. Is there a good way to do this by still using fn(x)? In reality, my function will be much more complicated - so I only want to run it once within the apply and assign output[0], output[1], etc to individual columns.

Grijesh Chauhan
  • 57,103
  • 20
  • 141
  • 208
arosner09
  • 649
  • 2
  • 8
  • 9
  • 6
    I think you can just write `df["y"], df["z"] = zip(*df["x"].apply(fn))`, as done in the answer to [this question](http://stackoverflow.com/questions/12356501/pandas-create-two-new-columns-in-a-dataframe-with-values-calculated-from-a-pre/). – DSM Sep 08 '14 at 16:51
  • Perfect - I think this works actually. Much appreciated! – arosner09 Sep 08 '14 at 17:12

1 Answers1

3

How about this method? (n.b., I edited this answer in light of the comment below) so the apply step could take a single function with shared calculations and return the required series for the merge step.

data = {'state':['Ohio','Ohio','Ohio','Nevada','Nevada'], 'year':[2000,2001,2002,2001,2002],'pop':[1.5,1.7,3.6,2.4,2.9]}
frame = pd.DataFrame(data, columns = ['year','state','pop'])
def fn(x,head1,head2):
    y = x ** 2
    z = x ** 3
    return pd.Series({head1:y, head2:z}) 
frame = frame.merge(frame['pop'].apply(lambda s: fn(s,'xsqr','xcube')), left_index=True, right_index=True)

Results:

   year   state  pop   xcube   xsqr
0  2000    Ohio  1.5   3.375   2.25
1  2001    Ohio  1.7   4.913   2.89
2  2002    Ohio  3.6  46.656  12.96
3  2001  Nevada  2.4  13.824   5.76
4  2002  Nevada  2.9  24.389   8.41
BKay
  • 1,397
  • 1
  • 15
  • 26
  • Thanks BKay - the issue is I need to use fn(x) and not individual calculations, e.g. s**2 and s**3. fn(x) is pretty complex, so I can't isolate the computation for y and z respectively. – arosner09 Sep 08 '14 at 17:02