Pandas - Cumsum，如果条件基于结果累加值，则跳过行

Question

How to accumulate values skipping rows if the accumulated result of those rows exceeds a certain threshold?

threshold = 120
Col1
---
100
5
90
5
8

Expected output:
Acumm_with_condition
---
100
105     (100+5)
NaN     (105+90 > threshold, skip )
110     (105+5)
118     (110+8)

Answer 1

Though it's not entirely vectorized, you can use a loop where you calculate the cumsum, then check to see if it has exceeded the threshold and if it has, set the value where it first breaks the threshold to 0 and restart the loop.

def thresholded_cumsum(df, column, threshold=np.inf, dropped_value_fill=None):
    s = df[column].copy().to_numpy()
    dropped_value_mask = np.zeros_like(s, dtype=bool)
    
    cur_cumsum = s.cumsum()
    cur_mask = cur_cumsum > threshold
    
    while cur_mask.any():
        first_above_thresh_idx = np.nonzero(cur_mask)[0][0]
        
        # Drop the value out of s, note the position of this value within the mask
        s[first_above_thresh_idx] = 0
        dropped_value_mask[first_above_thresh_idx] = True

        # Recalculate the cumsum & threshold mask now that we've dropped the value
        cur_cumsum = s.cumsum()
        cur_mask = cur_cumsum > threshold
            
    if dropped_value_fill is not None:
        cur_cumsum[dropped_value_mask] = dropped_value_fill
        
    return cur_cumsum

Usage:

df["thresh_cumsum"] = thresholded_cumsum(df, "col1", threshold=120)

print(df)
   col1  thresh_cumsum
0   100            100
1     5            105
2    90            105
3     5            110
4     8            118

I've included an extra parameter here dropped_value_fill , this is essentially a value you can use to annotate your output to let you know which values were intentionally dropped for violating the threshold.

With dropped_value_fill=-1

df["thresh_cumsum"] = thresholded_cumsum(df, "col1", threshold=120, dropped_value_fill=-1)

print(df)
   col1  thresh_cumsum
0   100            100
1     5            105
2    90             -1
3     5            110
4     8            118

Answer 2

Ended up using:

    def accumulate_under_threshold(values, threshold, skipped_row_value):
        output = []
        accumulated = 0
        for i, val in enumerate(values):
            if val + accumulated <= threshold:
                accumulated = val + accumulated
                output.append(accumulated)
            else:
                output.append(math.nan)
                if values[i:].min() > (threshold - accumulated ):
                    output.extend( [skipped_row_value]*(len(values)-1-i))
                    break
        return np.array(output)
    
    df['acumm_with_condition'] = accumulate_under_threshold(df['Col1'].values, 120, math.nan)