Início / coding / Perguntas / 79437667
Accepted
Simon
Asked: 2025-02-14 04:40:19 +0800 CST

Como contar combinações de estados exclusivos por ID em um Polars DataFrame

  • 772

Tenho um Polars DataFrame onde cada id pode aparecer várias vezes com valores de estado diferentes (1 ou 2). Quero contar quantos ids exclusivos têm apenas o estado 1, apenas o estado 2 ou ambos os estados 1 e 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]
})

Quero contar quantos IDs exclusivos se enquadram em cada categoria:

• Apenas indique 1 (por exemplo, IDs que têm apenas 1)

• Apenas indique 2 (por exemplo, IDs que tenham apenas 2)

• Ambos os estados 1 e 2 (por exemplo, IDs que têm 1 e 2)

Resultado esperado (exemplo):

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

3 respostas

  1. Você precisa de dois group_bys, o primeiro para combinar os ids nos estados e depois outro group_bypara os estados para contar o número de ids

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

    Você pode agrupar por ide usar .all()e .any()para verificar os estados.

    (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  │
└─────┴───────┴───────┴───────┘

    As .sum()colunas bool são as contagens.

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

  3. Para referência futura, forneça qual deve ser a saída correta em vez de apenas um exemplo de saída.

    Você pode executar um grupo, pegar os estados exclusivos para cada ID e, em seguida, pegar as contagens de valor disso

    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

relate perguntas