How to optimize conditional cumsum of multiple column

user_8481984 0
Viewed 88

I have very similar dataframe as below:

      Date        Col1    Col2    Col3 ...
   2020/11/04     -10      0       12
   2020/11/05      31      12      42
   2020/11/07      10      1      -12
   2020/11/08      2      -15      1
   2020/11/09      2       10      0
   .
   .
   .

My cumsum condition is while calculating sum for next row if sum is negative change it to 0. Output of the operation should look like below.

      Date        Col1    Col2    Col3 ...
   2020/11/04       0      0       12
   2020/11/05      31      12      54
   2020/11/07      41      13      42
   2020/11/08      43      0       43
   2020/11/09      45      10      43
   .
   .
   .

I have achieved this by looping and applying condition through each rows and column but for cumbersome data its performance is very poor.

columns = diff.columns
    for col in columns:
        if diff.iloc[0].at[col] < 0:
            diff.iloc[0].at[col] = 0
    for i,row in diff.iterrows():
        if not i == diff.first_valid_index():
            prev = diff.index.get_loc(i) - 1
            for col in columns:
                diff.loc[i].at[col] = diff.loc[i].at[col] + diff.iloc[prev].at[col]
                if diff.loc[i].at[col] < 0:
                    diff.loc[i].at[col] = 0

How can I do it better way in pandas?

UPDATE this thread here is very relevant and my solution is:

def adj_func(x):
    total = 0
    result = []
    for i, y in enumerate(x):
        total += y
        if total < 0:
            total = 0
        result.append(total)
    return result
diff[col_list].apply(adj_func)
pandas nested-loops python
1 Answers

njit

If you have access to Numba

from numba import njit

@njit
def f(a):
    for i, x in enumerate(a):
        if x < 0:
            a[i:] -= x
    return a

def g(s):
    return f(s.to_numpy().cumsum())

df.assign(**df[['Col1', 'Col2', 'Col3']].apply(g))

         Date  Col1  Col2  Col3
0  2020/11/04     0     0    12
1  2020/11/05    31    12    54
2  2020/11/07    41    13    42
3  2020/11/08    43     0    43
4  2020/11/09    45    10    43

Otherwise, without Numba

def h(s):
    a = s.to_numpy().cumsum()
    for i, x in enumerate(a):
        if x < 0:
            a[i:] -= x
    return a

df.assign(**df[['Col1', 'Col2', 'Col3']].apply(h))
user_23366540
Related