lassence.net

Field notes from the AI frontier

Get first row(s) of each group

--- archive: codebook, 2020–2023 · topic: Pandas ---

To get the first row, or first n rows, of each group in a pandas DataFrame, you can use groupby() followed by head().

If your DataFrame is not sorted, make sure you sort it before grouping.

And after the operation, you may want to reset indexes.

Here is an example:

# Import libraries
import pandas as pd
import seaborn as sns

# Load sample data in a DataFrame
df = (
    sns.load_dataset('iris')
    .sample(n=12, random_state=24)
    .sort_values('species')
    .reset_index(drop=True)
    [['species', 'sepal_width']]
)
df
speciessepal_width
0setosa3.4
1setosa3.7
2setosa3.2
3setosa3.1
4setosa3.8
5versicolor2.4
6versicolor2.7
7versicolor2.5
8virginica2.9
9virginica2.8
10virginica3.0
11virginica3.4
# Get smallest value of each group
(
    df
    .sort_values('sepal_width')
    .groupby('species')
    .head(1)
    .reset_index(drop=True)
)
speciessepal_width
0versicolor2.4
1virginica2.8
2setosa3.1

And that’s it! A couple of last remarks:

  • If you want to get the last row, just sort by inverse order with ascending=True
  • And to get the nth row, use nth() instead of head()