13  Grouping and Pivot Tables

In pandas, two common tools for summarizing data are:

import pandas as pd

df = pd.DataFrame({
    "state": ["CDMX", "CDMX", "Jalisco", "Jalisco", "Nuevo León", "Nuevo León"],
    "sex": ["M", "F", "M", "F", "M", "F"],
    "sector": ["Public", "Public", "Private", "Public", "Private", "Private"],
    "employees": [120, 150, 80, 110, 200, 170],
    "salary": [25000, 28000, 22000, 24000, 32000, 30000]
})

df
state sex sector employees salary
0 CDMX M Public 120 25000
1 CDMX F Public 150 28000
2 Jalisco M Private 80 22000
3 Jalisco F Public 110 24000
4 Nuevo León M Private 200 32000
5 Nuevo León F Private 170 30000

13.1 1. Basic groupby()

General structure:

df.groupby("group_variable")["numeric_variable"].aggregation()

13.1.1 Sum

df.groupby("state")["employees"].sum()
state
CDMX          270
Jalisco       190
Nuevo León    370
Name: employees, dtype: int64

This returns a Series.

To return a DataFrame:

df.groupby("state", as_index=False)["employees"].sum()
state employees
0 CDMX 270
1 Jalisco 190
2 Nuevo León 370

13.2 2. Count observations

13.2.1 size()

size() counts rows in each group.

df.groupby("state").size()
state
CDMX          2
Jalisco       2
Nuevo León    2
dtype: int64

A convenient DataFrame version:

df.groupby("state").size().reset_index(name="n")
state n
0 CDMX 2
1 Jalisco 2
2 Nuevo León 2

13.2.2 count()

count() counts non-missing values.

df.groupby("state")["salary"].count()
state
CDMX          2
Jalisco       2
Nuevo León    2
Name: salary, dtype: int64

The difference is:

size()  -> counts rows
count() -> counts non-null values

13.3 3. Group using more than one variable

df.groupby(["state", "sex"])["employees"].sum()
state       sex
CDMX        F      150
            M      120
Jalisco     F      110
            M       80
Nuevo León  F      170
            M      200
Name: employees, dtype: int64

Usually it is easier to work with the result using:

df.groupby(
    ["state", "sex"],
    as_index=False
)["employees"].sum()
state sex employees
0 CDMX F 150
1 CDMX M 120
2 Jalisco F 110
3 Jalisco M 80
4 Nuevo León F 170
5 Nuevo León M 200

13.4 4. Aggregation with agg()

agg() allows us to calculate several statistics at the same time.

df.groupby("state")["salary"].agg([
    "count",
    "sum",
    "mean",
    "min",
    "max"
])
count sum mean min max
state
CDMX 2 53000 26500.0 25000 28000
Jalisco 2 46000 23000.0 22000 24000
Nuevo León 2 62000 31000.0 30000 32000

You can also aggregate several numeric variables:

df.groupby("state").agg({
    "employees": "sum",
    "salary": "mean"
})
employees salary
state
CDMX 270 26500.0
Jalisco 190 23000.0
Nuevo León 370 31000.0

13.5 5. Named aggregations

Named aggregations are particularly useful because we can control the names of the resulting columns.

summary = (
    df
    .groupby("state", as_index=False)
    .agg(
        n=("salary", "size"),
        total_employees=("employees", "sum"),
        average_salary=("salary", "mean"),
        minimum_salary=("salary", "min"),
        maximum_salary=("salary", "max")
    )
)

summary
state n total_employees average_salary minimum_salary maximum_salary
0 CDMX 2 270 26500.0 25000 28000
1 Jalisco 2 190 23000.0 22000 24000
2 Nuevo León 2 370 31000.0 30000 32000

General syntax:

new_column=("original_column", "aggregation")

For example:

total=("employees", "sum")

13.6 6. Group categorical variables

Categorical variables can also be summarized.

For example, count observations by sector:

df.groupby("sector").size()
sector
Private    3
Public     3
dtype: int64

Or:

df.groupby("sector").size().reset_index(name="n")
sector n
0 Private 3
1 Public 3

Count observations for combinations of categories:

df.groupby(["state", "sector"]).size().reset_index(name="n")
state sector n
0 CDMX Public 2
1 Jalisco Private 1
2 Jalisco Public 1
3 Nuevo León Private 2

13.7 7. Count unique categories with nunique()

Suppose we want to know how many different sectors appear in each state.

df.groupby("state")["sector"].nunique()
state
CDMX          1
Jalisco       2
Nuevo León    1
Name: sector, dtype: int64

Or:

df.groupby("state", as_index=False).agg(
    sectors=("sector", "nunique")
)
state sectors
0 CDMX 1
1 Jalisco 2
2 Nuevo León 1

14 Pivot tables

pivot_table() works similarly to a pivot table in Excel.

General structure:

pd.pivot_table(
    df,
    index="row_variable",
    columns="column_variable",
    values="numeric_variable",
    aggfunc="aggregation"
)

14.1 8. Sum in a pivot table

Sum employees by state and sex:

pd.pivot_table(
    df,
    index="state",
    columns="sex",
    values="employees",
    aggfunc="sum"
)
sex F M
state
CDMX 150 120
Jalisco 110 80
Nuevo León 170 200

14.2 9. Mean in a pivot table

Average salary by state and sex:

pd.pivot_table(
    df,
    index="state",
    columns="sex",
    values="salary",
    aggfunc="mean"
)
sex F M
state
CDMX 28000.0 25000.0
Jalisco 24000.0 22000.0
Nuevo León 30000.0 32000.0

14.3 10. Count observations in a pivot table

One way to count observations is:

pd.pivot_table(
    df,
    index="state",
    columns="sex",
    values="salary",
    aggfunc="count"
)
sex F M
state
CDMX 1 1
Jalisco 1 1
Nuevo León 1 1

However, for purely categorical counts, pd.crosstab() is often simpler:

pd.crosstab(
    df["state"],
    df["sex"]
)
sex F M
state
CDMX 1 1
Jalisco 1 1
Nuevo León 1 1

14.4 11. Multiple aggregation functions

A pivot table can calculate several statistics simultaneously.

pd.pivot_table(
    df,
    index="state",
    columns="sex",
    values="salary",
    aggfunc=["count", "mean", "sum"]
)
count mean sum
sex F M F M F M
state
CDMX 1 1 28000.0 25000.0 28000 25000
Jalisco 1 1 24000.0 22000.0 24000 22000
Nuevo León 1 1 30000.0 32000.0 30000 32000

14.5 12. Multiple numeric variables

pd.pivot_table(
    df,
    index="state",
    values=["employees", "salary"],
    aggfunc="sum"
)
employees salary
state
CDMX 270 53000
Jalisco 190 46000
Nuevo León 370 62000

Different aggregations can also be assigned to different variables:

pd.pivot_table(
    df,
    index="state",
    values=["employees", "salary"],
    aggfunc={
        "employees": "sum",
        "salary": "mean"
    }
)
employees salary
state
CDMX 270 26500.0
Jalisco 190 23000.0
Nuevo León 370 31000.0

14.6 13. Multiple grouping variables

pd.pivot_table(
    df,
    index=["state", "sector"],
    columns="sex",
    values="employees",
    aggfunc="sum"
)
sex F M
state sector
CDMX Public 150.0 120.0
Jalisco Private NaN 80.0
Public 110.0 NaN
Nuevo León Private 170.0 200.0

14.7 14. Replace missing combinations with zero

Pivot tables can generate missing values when a combination does not exist.

Use fill_value:

pd.pivot_table(
    df,
    index="state",
    columns="sector",
    values="employees",
    aggfunc="sum",
    fill_value=0
)
sector Private Public
state
CDMX 0 270
Jalisco 80 110
Nuevo León 370 0

14.8 15. Add totals

Use margins=True:

pd.pivot_table(
    df,
    index="state",
    columns="sex",
    values="employees",
    aggfunc="sum",
    margins=True,
    margins_name="Total"
)
sex F M Total
state
CDMX 150 120 270
Jalisco 110 80 190
Nuevo León 170 200 370
Total 430 400 830

15 Quick reference

15.0.1 Count rows

df.groupby("category").size()

15.0.2 Count non-null values

df.groupby("category")["variable"].count()

15.0.3 Sum

df.groupby("category")["numeric"].sum()

15.0.4 Mean

df.groupby("category")["numeric"].mean()

15.0.5 Multiple statistics

df.groupby("category")["numeric"].agg([
    "count",
    "sum",
    "mean",
    "min",
    "max"
])

15.0.7 Pivot table

pd.pivot_table(
    df,
    index="category_1",
    columns="category_2",
    values="numeric",
    aggfunc="sum"
)

15.0.8 Categorical frequency table

pd.crosstab(
    df["category_1"],
    df["category_2"]
)