pandas groupby expanding df based on unique values

Viewed 673

I have df below:

df = pd.DataFrame({
    'ID': ['a', 'a', 'a', 'b', 'c', 'c'],
    'V1': [False, False, True, True, False, True],
    'V2': ['A', 'B', 'C', 'B', 'B', 'C']
})

I want to achieve the following. For each unique ID, the bottom row is True (this is V1). I want to count how many times each unique value of V2 occurs where V1==True. This part would be achieved by something like:

df.groupby('V2').V1.sum()

However, I also want to add, for each unique value of V2, a column indicating how many times that value occurred after the point where V1==True for the V2 value indicated by the row. I understand this might sound confusing; here's how the output woud look like in this example:

df
    V2  V1  A   B   C
0   A   0   0   0   0
1   B   1   0   0   0
2   C   2   1   2   0

It is important that the solution is general enough to be applicable on a similar case with more unique values than just A, B and C.

UPDATE
As a bonus, I am also interested in how, instead of the count, one can instead return the sum of some value column, under the same conditions, divided by the corresponding "count" in the rows. Example: suppose we now depart from df below instead:

df = pd.DataFrame({
    'ID': ['a', 'a', 'a', 'b', 'c', 'c'],
    'V1': [False, False, True, True, False, True],
    'V2': ['A', 'B', 'C', 'B', 'B', 'C'],
    'V3': [1, 2, 3, 4, 5, 6],
})

The output would need to sum V3 for the cases indicated by the counts in the solution by @jezrael, and divide that number by V1. The output would instead look like:

df
    V2  V1  A   B   C
0   A   0   0   0   0
1   B   1   0   0   0
2   C   2   1   3.5 0
1 Answers

First aggregate sum:

df1 = df.groupby('V2').V1.sum().astype(int).reset_index()
print (df1)
  V2  V1
0  A   0
1  B   1
2  C   2

Then grouping by ID and create heper column by last value by GroupBy.transform and last, then remove last rows of ID by Series.duplicated and use crosstab for counts, add all possible unique values of V2 and last append to df1 by DataFrame.join:

val = df['V2'].unique()
df['new'] = df.groupby('ID').V2.transform('last')
df = df[df.duplicated('ID', keep='last')]

df = pd.crosstab(df['new'], df['V2']).reindex(columns=val, index=val, fill_value=0)

df = df1.join(df, on='V2')
print (df)
  V2  V1  A  B  C
0  A   0  0  0  0
1  B   1  0  0  0
2  C   2  1  2  0

UPDATE

The updated part of the question should be possible to achieve by changing the crosstab part with pivot table:

df = df.pivot_table(
    index='n',
    columns='V2',
    aggfunc=({
        'V3': 'mean'
    })
).V3.reindex(columns=v, index=v, fill_value=0)
Related