How to do a running count based on certain conditions in Python

Viewed 70

I currently have this Dataset:

stack = [('2021-04-25', 9, 0, 0, 0, 0, 'A'),
           ('2021-04-26', 0, 0, 0, 0, 0, 'A'),
           ('2021-04-27', 0, 0, 0, 0, 9, 'A'),
           ('2021-04-28', 56, 0, 0, 0, 53, 'A'),
           ('2021-04-29', 1, 0, 1, 0, 0, 'A'),
           ('2021-04-30', 3, 0, 0, 0, 4, 'A'),
           ('2021-05-01', 0, 0, 0, 0, 0, 'A'),
           ('2021-05-02', 21, 0, 0, 0, 0, 'A'),
           ('2021-05-03', 4, 0, 0, 0, 0, 'A'),
           ('2021-05-04', 1, 0, 0, 0, 0, 'A')]
stack = pd.DataFrame(stack, columns=['Date', 'V', 'W', 'X', 'Y', 'Z', 'Group'])

And would like to get this as my final output:

stack_out = [('2021-04-25', 9, 0, 0, 0, 0, 'A'),
           ('2021-04-26', 9, 0, 0, 0, 0, 'A'),
           ('2021-04-27', 9, 0, 0, 0, 9, 'A'),
           ('2021-04-28', 56, 0, 0, 0, 62, 'A'),
           ('2021-04-29', 4, 0, 1, 0, 62, 'A'),
           ('2021-04-30', 6, 0, 0, 0, 66, 'A'),
           ('2021-05-01', 0, 0, 0, 0, 68, 'A'),
           ('2021-05-02', 21, 0, 0, 0, 68, 'A'),
           ('2021-05-03', 25, 0, 0, 0, 68, 'A'),
           ('2021-05-04', 26, 0, 0, 0, 68, 'A')]
stack_out = pd.DataFrame(stack_out, columns=['Date', 'V', 'W', 'X', 'Y', 'Z', 'Group'])

A few conditions that are taking place:

  1. Anytime anything occurs in W, X, Y, Z it must be subtracted from V.
  2. In stack_out, I would like V to be a running total of how many remain after subtraction.
  3. At the end of every month, the count in V resets to 0 and any number that remains in V at the end of the month should be added to Z at the beginning of the month.
  4. Z should be a running total as well + the number added to it at the end of every month.

Clarification:

  1. On 4/28 there are 56 in V and 53 in Z. Subtracting these should lead to a reminder of 3 and that should be added to the next day. That is why there is a 4 in 4/29 in stack_out.
  2. In 4/29 the 4 is subtracted from the 1 in X and the 3 remainder are carried over to the next day leading to 6 on 4/30.
  3. On 4/30, the 6 are subtracted from the 4 in Z leading to a remainder of 2. Since it is the end of the month, V 5/01 will reset to 0 and the remaining 2 will be added to Z hence why Z increased from 66 to 68 on 5/01.
1 Answers

This isn't the prettiest code and could be improved, but I think this will work for you. I did it in two steps, first calculating the cumulative sum and then subtracting/adding the values on the first day of each month to each value for the month in V/Z.

Edit: I vectorized the method using groupby instead of using iterrows.

stack = [('2021-04-25', 9, 0, 0, 0, 0, 'A'),
           ('2021-04-26', 0, 0, 0, 0, 0, 'A'),
           ('2021-04-27', 0, 0, 0, 0, 9, 'A'),
           ('2021-04-28', 56, 0, 0, 0, 53, 'A'),
           ('2021-04-29', 1, 0, 1, 0, 0, 'A'),
           ('2021-04-30', 3, 0, 0, 0, 4, 'A'),
           ('2021-05-01', 0, 0, 0, 0, 0, 'A'),
           ('2021-05-02', 21, 0, 0, 0, 0, 'A'),
           ('2021-05-03', 4, 0, 0, 0, 0, 'A'),
           ('2021-05-04', 1, 0, 0, 0, 0, 'A'),]
df = pd.DataFrame(stack, columns=['Date', 'V', 'W', 'X', 'Y', 'Z', 'Group'])
df['Date'] = pd.to_datetime(df['Date'])

# Create Cumulative Sums for V and Z
df['V'] = (df['V'] - (
    df[['W', 'X', 'Y', 'Z']]
    .sum(axis=1)
    .shift(fill_value=0)
)).cumsum()

df['Z'] = df['Z'].cumsum()

# Group by month and add the first of month's value to all values in Z 
# and subtract from V
df['Z'] = df.groupby(df['Date'].dt.month)['Z'].transform(lambda x: (
    x + (
        df.loc[(df['Date'].dt.month == x.name) & (df['Date'].dt.is_month_start), 'V'].values
        if len(df.loc[(df['Date'].dt.month == x.name) & (df['Date'].dt.is_month_start), 'V'].values) > 0
        else 0
    )
))
df['V'] = df.groupby(df['Date'].dt.month)['V'].transform(lambda x: (
    x - (
        df.loc[(df['Date'].dt.month == x.name) & (df['Date'].dt.is_month_start), 'V'].values
        if len(df.loc[(df['Date'].dt.month == x.name) & (df['Date'].dt.is_month_start), 'V'].values) > 0
        else 0
    )
))

Output

        Date   V  W  X  Y   Z Group
0 2021-04-25   9  0  0  0   0     A
1 2021-04-26   9  0  0  0   0     A
2 2021-04-27   9  0  0  0   9     A
3 2021-04-28  56  0  0  0  62     A
4 2021-04-29   4  0  1  0  62     A
5 2021-04-30   6  0  0  0  66     A
6 2021-05-01   0  0  0  0  68     A
7 2021-05-02  21  0  0  0  68     A
8 2021-05-03  25  0  0  0  68     A
9 2021-05-04  26  0  0  0  68     A
Related