0

I am trying to convert VBA macro to python script. I need to add a new column to dataframe:

 X  Y 
    
 1  2 
 3  4
 5  6 
 7  8

that will be a sum of a value above (0 for the first one) and value from X multiplied by Y summed with a value above :

 X  Y  Z    Z(calculations) 

 1  2  0    0 
 3  4  12   0 + 3*(4+0)
 5  6  102  12 + 5*(6+12)  
 7  8  872  102 + 7*(8+102)

So basically Excel formula (in Z2 Cell):

=Z1+X2*(Y2+Z1)

Is there an easy (and quick) way to do this?

Pandriej
  • 21
  • 1
  • 3

1 Answers1

0

You could do it with a for loop but since you asked for a quick way you can do this using numba and @njit:

from numba import njit

import pandas as pd
import numpy as np

@njit
def calculate(X, Y):
    res = np.empty(X.shape)
    res[0] = 0
    for i in range(1, res.shape[0]):
        res[i] = res[i-1] + (X[i] * (Y[i] + res[i-1]))
    return res

df = pd.DataFrame({"X": [1,3,5,7], "Y": [2,4,6,8]})
df
   X  Y
0  1  2
1  3  4
2  5  6
3  7  8

df["Z"] = calculator_nb(*df[["X", "Y"]].values.T)

df
   X  Y      Z
0  1  2    0.0
1  3  4   12.0
2  5  6  102.0
3  7  8  872.0

Credit to this post

Benjamin Carlier
  • 77
  • 6