主页 / coding / 问题 / 79437667
Accepted
Simon
Asked: 2025-02-14 04:40:19 +0800 CST

如何计算 Polars DataFrame 中每个 ID 的唯一状态组合

  • 772

我有一个 Polars DataFrame，其中每个 id 可以出现多次，并且具有不同的状态值（1 或 2）。我想计算有多少个唯一 id 只具有状态 1、只具有状态 2 或同时具有状态 1 和 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]
})

我想计算每个类别中有多少个唯一的 ID：

• 仅状态 1（例如，只有 1 的 ID）

• 仅陈述 2（例如，只有 2 个的 ID）

• 状态 1 和 2（例如，同时具有 1 和 2 的 ID）

预期结果（示例）：

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

3 个回答

  1. 您需要两个group_bys，第一个 s 将ids 合并到状态中，另一个sgroup_by用于统计状态的数量id

    (
    df
    .group_by("id")
    .agg(pl.col("state").unique().sort())
    .group_by("state")
    .len()
)
    • 3
  2. Best Answer

    您可以按分组id并使用.all()和.any()检查状态。

    (df.group_by("id")
   .agg(
       one  = (pl.col.state == 1).all(),
       two  = (pl.col.state == 2).all(),
       both = (pl.col.state == 1).any() & (pl.col.state == 2).any()
       # both = pl.lit(1).is_in("state") & pl.lit(2).is_in("state")
   )
   # .select(pl.exclude("id").sum())
)

    shape: (20, 4)
┌─────┬───────┬───────┬───────┐
│ id  ┆ one   ┆ two   ┆ both  │
│ --- ┆ ---   ┆ ---   ┆ ---   │
│ i64 ┆ bool  ┆ bool  ┆ bool  │
╞═════╪═══════╪═══════╪═══════╡
│ 6   ┆ false ┆ true  ┆ false │
│ 3   ┆ false ┆ true  ┆ false │
│ 2   ┆ true  ┆ false ┆ false │
│ 12  ┆ true  ┆ false ┆ false │
│ 16  ┆ false ┆ false ┆ true  │
│ …   ┆ …     ┆ …     ┆ …     │
│ 9   ┆ false ┆ false ┆ true  │
│ 13  ┆ false ┆ true  ┆ false │
│ 8   ┆ false ┆ true  ┆ false │
│ 15  ┆ false ┆ false ┆ true  │
│ 10  ┆ false ┆ false ┆ true  │
└─────┴───────┴───────┴───────┘

    .sum()布尔列是计数。

    shape: (1, 3)
┌─────┬─────┬──────┐
│ one ┆ two ┆ both │
│ --- ┆ --- ┆ ---  │
│ u32 ┆ u32 ┆ u32  │
╞═════╪═════╪══════╡
│ 6   ┆ 7   ┆ 7    │
└─────┴─────┴──────┘
    • 3

  3. 为了将来参考，请提供正确的输出，而不仅仅是示例输出。

    您可以执行分组，获取每个 ID 对应的唯一州，然后获取该州的值计数

    combinations = df.group_by('id').agg(pl.col('state').unique())
counts = combinations.select(pl.col('state').value_counts().alias('counts'))
print(counts.unnest('counts'))

assert (counts.select(pl.col('counts').struct.field('count').sum()) == df.n_unique('id')).item()

# Alternatively, as a single expression:
print(df.select(
    pl.col('state').unique().implode()
    .over('id', mapping_strategy='explode')
    .value_counts()
    .struct.unnest()
))
    • 2

相关问题