CODEWITHSUNDEEP
pythonpandasdataframenumpy
Marra
Marra

Reputation: 33

Calculate cumulative sum based on threshold and condition in another column numpy

I have a data frame and I'd like to calculate cumulative sum based on 2 conditions:

I've succeed with 1st or 2nd but I find it hard to combine both.

For the first one I used groupby

df['group'] = np.cumsum((df['IsSuccess'] != df['IsSuccess'].shift(1)))
df['SumSale'] = df[['Sale', 'group']].groupby('group').cumsum()

For the 2nd frompyfunc

sumlm = np.frompyfunc(lambda a,b: b if (a+b>5) else a+b, 2, 1)
df['SumSale'] = sumlm.accumulate(df['Sale'], dtype=object)

My df is, and the SumSale is the result I'm looking for.

df2 = pd.DataFrame({'Sale': [10, 2, 2, 1, 3, 2, 1, 3, 5, 5],
                 'IsSuccess': [False, True, False, False, True, False, True, False, False, False],
                 'SumSaleExpected': [10, 12, 2, 3, 6, 2, 3, 6, 11, 16]})

So to summarize I'd like to start cumulating the sum once that sum is over 5 and the row IsSuccess is True. I'd like to avoid for loop if possible as well.

Thank you for help!

Upvotes: 3

Views: 461

Answers (1)

Andrej Kesely
Andrej Kesely

Reputation: 195418

I hope I've understood your question right. This example will substract necessary value ("reset") when cumulative sum of sale is greater than 5 and IsSuccess==True:

df["SumSale"] = df["Sale"].cumsum()

# "reset" when SumSale>5 and IsSuccess==True
m = df["SumSale"].gt(5) & df["IsSuccess"].eq(True)
df.loc[m, "to_remove"] = df["SumSale"]
df["to_remove"] = df["to_remove"].ffill().shift().fillna(0)
df["SumSale"] -= df["to_remove"]

df = df.drop(columns="to_remove")

print(df)

Prints:

   Sale  IsSuccess  SumSale
0     1      False      1.0
1     2       True      3.0
2     3      False      6.0
3     2      False      8.0
4     4       True     12.0
5     3      False      3.0
6     5       True      8.0
7     5      False      5.0

EDIT:

def fn():
    sale, success = yield
    cum = sale
    while True:
        sale, success = yield cum
        if success and cum > 5:
            cum = sale
        else:
            cum += sale


s = fn()
next(s)
df["ss"] = df["IsSuccess"].shift()
df["SumSale"] = df.apply(lambda x: s.send((x["Sale"], x["ss"])), axis=1)
df = df.drop(columns="ss")
print(df)

Prints:

   Sale  IsSuccess  SumSaleExpected  SumSale
0    10      False               10       10
1     2       True               12       12
2     2      False                2        2
3     1      False                3        3
4     3       True                6        6
5     2      False                2        2
6     1       True                3        3
7     3      False                6        6
8     5      False               11       11
9     5      False               16       16

Upvotes: 1

Related Questions