Calculate cumilative sum for all department column

Question

I have this Dataframe

lst = [['AA','Z',10,1,0],['BB','Y',10,1,0],
       ['AA','Z',20,2,0],['CC','X',10,2,0]]
df1 = pd.DataFrame(lst,columns = ['first_name','last_name','val','department','is_cum'])

looks like this

  first_name last_name  val  department  is_cum
0         AA         Z   10           1      NO
1         BB         Y   10           1      NO
2         AA         Z   20           2      NO
3         CC         X   10           2      NO

I want output something like this

  first_name last_name  val  department  is_cum
0         AA         Z   10           1      NO
1         BB         Y   10           1      NO
2         AA         Z   10           1     YES
3         BB         Y   10           1     YES

4         AA         Z   20           2      NO
5         CC         X   10           2      NO
6         AA         Z   30         1,2     YES
7         CC         X   10           2     YES
8         BB         Y   10           1     YES

All the rows with is_cum NO is same as the input dataframe the newly populated rows are the cumulative rows with is_cum as YES.

Row 2 and 3 are the same as 0 and 1 as we have just one department to do cumulation. Row 6,7,8 is the cumulation of department 1 and department 2. If we have same first_name and last_name in department 1 and department 2 than add there val or else keep them as it is.

I was doing

df1.groupby(['first_name','last_name','department']).sum().groupby(level=0).cumsum()

after this, i could change the is_cum col and append these row in the original dataframe. But this is not the required output.

user_9274732 · Answer

Here is one way using pivot_table to be able to perform the cumsum along the columns. All the rest is pretty much getting the expected output.

df_ = (df1.assign(dpt=df1['department'], 
                  department=df1['department'].astype(str))\
          .pivot_table(index=['first_name','last_name'], columns='dpt', 
                       values=['val', 'department'], 
                       aggfunc={'val':sum,'department':lambda x: list(x)})
          .assign(val=lambda x: x['val'].cumsum(axis=1).ffill(axis=1), 
                  department=lambda x: x['department'].apply(lambda x: x.dropna().cumsum(), axis=1)
                                                        .ffill(axis=1))
                 )

res= (pd.concat([df1.assign(is_cum='NO', dpt=df1['department']), 
                 df_.stack().reset_index()
                    .assign(is_cum='YES',
                            department=lambda x: x['department'].apply(','.join))])
        .sort_values(['dpt', 'is_cum']).drop('dpt',axis=1)
        .reset_index(drop=True)
     )

and you get

print(res)
  first_name last_name   val department is_cum
0         AA         Z  10.0          1     NO
1         BB         Y  10.0          1     NO
2         AA         Z  10.0          1    YES
3         BB         Y  10.0          1    YES
4         AA         Z  20.0          2     NO
5         CC         X  10.0          2     NO
6         AA         Z  30.0        1,2    YES
7         BB         Y  10.0          1    YES
8         CC         X  10.0          2    YES

Calculate cumilative sum for all department column

1 Answers