Pandas: cumsum with conditions

Question

I have the following dataframe in pandas:

   A     B
   5     0
   5     10
   5     0
   5     0
   5     50

I would like to calculate a column C as the cumulative difference between B and A but with at the condition that C cannot be lower than 0 and higher than 15. How could I obtain the following result with the cumsum function?

   A     B     C
   5     0     0
   5     10    5
   5     0     0
   5     0     0
   5     50    15

Answer 1

@郑远程 has the correct approach (this answer should be accepted), as the desired operation requires a loop and cannot be vectorized:

Here is a just variant using a function and numba to improve efficiency for those interested:

import numba as nb

@nb.jit(nopython=True)
def cumreset(s):
    out = []
    tot = 0
    for v in s:
        tot = max(0, min(tot+v, 15))
        out.append(tot)
    return out

df['C'] = cumreset(df['B'].sub(df['A']).to_numpy())

Output:

   A   B   C
0  5   0   0
1  5  10   5
2  5   0   0
3  5   0   0
4  5  50  15

Comparison of the python and numba loops:

Answer 2

Oops! My apologies for that oversight. The issue is occurring because the cumulative_diff function is returning a list within the expanding().apply() method, causing a type mismatch.

To fix this, we can change our approach slightly by using a loop to calculate the cumulative sum. Here's the corrected code:

import pandas as pd

df = pd.DataFrame({
    'A': [5, 5, 5, 5, 5],
    'B': [0, 10, 0, 0, 50]
})

# Calculate the difference between B and A
df['diff'] = df['B'] - df['A']

# Initialize a variable to hold the cumulative difference
cum_diff = 0

# Create an empty list to hold the results
results = []

# Loop through the differences, calculate the cumulative sum, and append to the results
for value in df['diff']:
    cum_diff += value
    cum_diff = max(0, cum_diff)  # Reset to 0 if cumulative difference is less than 0
    cum_diff = min(15, cum_diff) # Reset to 15 if cumulative difference is more than 15
    results.append(cum_diff)

# Assign the results to the new 'C' column
df['C'] = results

# Remove the temporary difference column
df.drop('diff', axis=1, inplace=True)

print(df)

This will yield the desired output:

   A   B   C
0  5   0   0
1  5  10   5
2  5   0   0
3  5   0   0
4  5  50  15

I hope this corrected version works for you, and I'm sorry again for the confusion.

Answer 3

I am not sure that your desired output corresponds to your text. I have followed your text above:

df['diff'] = df['B'] - df['A']
df['diff'] = np.where((df['diff']>0) & (df['diff']<15), df['diff'], 0 )
df['C'] = df['diff'].cumsum()
df

   A   B  diff  C
0  5   0     0  0
1  5  10     5  5
2  5   0     0  5
3  5   0     0  5
4  5  50     0  5

Answer 4

Try this:

n = df['B'].cumsum().sub(df['A'].cumsum())

df.assign(C = df['A'].groupby(n.lt(0).iloc[::-1].cumsum()).cumsum().rsub(df['B'].cumsum()).clip(0,15).where(n.ge(0),0))

Output:

   A   B   C
0  5   0   0
1  5  10   5
2  5   0   0
3  5   0   0
4  5  50  15