Pandas pivot table using custom conditions on the dataframe

Question

I want to make a pivot table based on custom conditions in the dataframe:

The dataframe looks like this:

>>> df = pd.DataFrame({"Area": ["A", "A", "B", "A", "C", "A", "D", "A"],
                       "City" : ["X", "Y", "Z", "P", "Q", "R", "S", "X"],
                       "Condition" : ["Good", "Bad", "Good", "Good", "Good", "Bad", "Good", "Good"], 
                       "Population" : [100,150,50,200,170,390,80,100]
                       "Pincode" : ["X1", "Y1", "Z1", "P1", "Q1", "R1", "S1", "X2"] })
>>> df
  Area City Condition   Population Pincode
 0    A    X      Good   100       X1
 1    A    Y       Bad   150       Y1
 2    B    Z      Good   50        Z1
 3    A    P      Good   200       P1
 4    C    Q      Good   170       Q1
 5    A    R       Bad   390       R1
 6    D    S      Good   80        S1
 7    A    X      Good   100       X2

Now I want to pivot the dataframe df in a manner such that I can see the unique count of cities against each area along with the corresponding count of "Good" cities and also the population of the area.

I expect an output like this:

Area  city_count  good_city_count   Population
A        4        2                 940
B        1        1                 50
C        1        1                 170
D        1        1                 80
All      7        5                 1240

I can give a dictionary to the aggfunc parameter but this doesn't give me the city count split between the good cities.

>>> city_count = df.pivot_table(index=["Area"],
                                values=["City", "Population"],
                                aggfunc={"City": lambda x: len(x.unique()),
                                         "Population": "sum"},
                                margins=True)

    Area    City    Population
0   A       4       940
1   B       1       50
2   C       1       170
3   D       1       80
4   All     7       1240

I can merge two different pivot tables - one with the count of cities and the other with the population but this is not scalable for a large dataset with a big aggfunc dictionary.

Answer 1

Another method without using pivot_table. Use np.where with groupby+agg:

df['Condition'] = np.where(df['Condition']=='Good', df['City'], np.nan)
df = df.groupby('Area').agg({'City':'nunique', 'Condition':'nunique', 'Population':'sum'})\
                       .rename(columns={'City':'city_count', 'Condition':'good_city_count'})
df.loc['All',:] = df.sum()
df = df.astype(int).reset_index()

print(df)
  Area  city_count  good_city_count  Population
0    A           4                2         940
1    B           1                1          50
2    C           1                1         170
3    D           1                1          80
4  All           7                5        1240

Answer 2

Add new parameters columns with fill_value and also is possible use nunique for aggregate function:

city_count = df.pivot_table(index = "Area", 
                            values = "City", 
                            columns='Condition', 
                            aggfunc = lambda x : x.nunique(), 
                            margins = True,
                            fill_value=0)
print (city_count)
Condition  Bad  Good  All
Area                     
A            2     2    4
B            0     1    1
C            0     1    1
D            0     1    1
All          2     5    7

Last if need convert index to column and change columns names:

city_count = city_count.add_suffix('_count').reset_index().rename_axis(None, 1)
print (city_count)
  Area  Bad_count  Good_count  All_count
0    A          2           2          4
1    B          0           1          1
2    C          0           1          1
3    D          0           1          1
4  All          2           5          7

EDIT:

d = {'City':'nunique','Population':'sum', 'good_city_count':'nunique'}
d1 = {'City':'city_count','Condition':'good_city_count'}

mask = df["Condition"] == 'Good'
df = (df.assign(good_city_count = lambda x: np.where(mask, x['City'], np.nan))
       .groupby('Area')
       .agg(d)
       .rename(columns=d1))

df = df.append(df.sum().rename('All')).reset_index()

print (df)
  Area  city_count  Population  good_city_count
0    A           4         940                2
1    B           1          50                1
2    C           1         170                1
3    D           1          80                1
4  All           7        1240                5