I have DataFrame as below:
df = pd.DataFrame((np.random.randn(5,4)*10).astype(int), columns=list('abcd'))
def cal(a, b):
if a + b > 5:
return a+b, a-b
how could I apply this function to df, the two variables cal take would be df['a'] and ['b'], the output a+b, a-b will be set to df['c'], df['d'].
loop the df works, but How could I use apply or applymap to achieve this (maybe cal need to be tweak)?
You can vectorize this using mask:
vals = pd.concat((df['a'] + df['b'], df['a'] - df['b']), axis=1).values
df[['c', 'd']].mask(df['a'] + df['b'] > 5, vals)
Out:
c d
0 6 3
1 -12 3
2 12 -14
3 21 -31
4 15 -21
where the original df is
df
Out:
a b c d
0 3 -2 6 3
1 -4 -8 -12 3
2 -1 13 10 -4
3 -5 26 -21 17
4 -3 18 14 19
np.random.seed([3,1415])
df = pd.DataFrame(
(np.random.randn(5, 4) * 10).astype(int),
columns=list('abcd')
)
df
a b c d
0 -21 -12 -19 -22
1 -3 0 3 3
2 7 7 -6 3
3 -3 -4 -9 -1
4 7 -15 6 4
Use dot and multiply the appropriate transformation
v = np.column_stack([df.a.values, df.b.values])
w = np.column_stack([df.c.values, df.d.values])
trans = np.array([[1, 1], [1, -1]])
df.loc[:, ['c', 'd']] = np.where(v.sum(1, keepdims=True) > 5, v.dot(trans), w)
df
a b c d
0 -21 -12 -19 -22
1 -3 0 3 3
2 7 7 14 0
3 -3 -4 -9 -1
4 7 -15 6 4
