Ndtyjky

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.

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

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