如何在Polars DataFrame中按ID统计唯一状态组合数量

需求：针对包含重复ID和状态值（仅1或2）的Polars DataFrame，统计以下三类唯一ID的数量：

• 仅出现状态1的ID
• 仅出现状态2的ID
• 同时出现状态1和2的ID

示例数据：

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]
})

解决方案代码 #

# 按ID分组提取唯一状态，再按状态组合统计ID数量
result = (
    df
    .group_by("id")
    .agg(pl.col("state").unique().sort().alias("state_combination"))
    .group_by("state_combination")
    .agg(pl.count("id").alias("id_count"))
)

# 格式化输出结果
for row in result.iter_rows(named=True):
    comb = row["state_combination"]
    count = row["id_count"]
    print(f"状态组合 {comb}  -> {count} 个ID")

代码说明 #

1. 按ID分组后，用unique().sort()提取每个ID对应的所有唯一状态并排序，避免同一组合出现不同顺序的格式（比如[1,2]和[2,1]被判定为不同组合）。
2. 再按整理好的状态组合分组，统计每组对应的唯一ID总数。
3. 最后遍历结果，按指定格式输出统计数据。

运行结果 #

状态组合 [1]  -> 5 个ID  
状态组合 [2]  -> 6 个ID  
状态组合 [1, 2]  -> 9 个ID  

内容的提问来源于stack exchange，提问作者Simon