[Python] How to count unique state combinations per ID in a Polars DataFrame

I have a Polars DataFrame where each id can appear multiple times with different state values (either 1 or 2). I want to count how many unique ids have only state 1, only state 2, or both states 1 and 2.

import polars as pl

df = pl.DataFrame({
"id": [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 9, 9, 10, 10, 10, 11, 11, 12, 12, 13, 14, 15, 15, 16, 16, 17, 17, 18, 18, 19, 20, 20, 20],
"state": [1, 2, 1, 1, 2, 2, 1, 2, 1, 1, 2, 2, 1, 1, 2, 1, 2, 1, 2, 2, 2, 2, 1, 1, 2, 2, 1, 2, 1, 2, 1, 1, 2, 2, 1, 1, 2, 2]
})


I want to count how many unique ids fall into each category:

• Only state 1 (e.g., IDs that only have 1)

• Only state 2 (e.g., IDs that only have 2)

• Both states 1 and 2 (e.g., IDs that have both 1 and 2)

Expected Result (Example):

State combination [1] -> 20 IDs
State combination [2] -> 15 IDs
State combination [1, 2] -> 30 IDs

Continue reading...