Início / coding / Perguntas / 78913464
Accepted
DJDuque
Asked: 2024-08-26 16:27:52 +0800 CST

Como filtrar exclusividade por condição

  • 772

Imagine que tenho um conjunto de dados como:

data = {
    "a": [1, 4, 2, 4, 7, 4],
    "b": [4, 2, 3, 3, 0, 2],
    "c": ["a", "b", "c", "d", "e", "f"],
}

e quero manter apenas as linhas a + bdescritas exclusivamente por uma única combinação de ae b. Eu consegui hackear isso:

df = (
    pl.DataFrame(data)
    .with_columns(sum_ab=pl.col("a") + pl.col("b"))
    .group_by("sum_ab")
    .agg(pl.col("a"), pl.col("b"), pl.col("c"))
    .filter(
        (pl.col("a").list.unique().list.len() == 1)
        & (pl.col("b").list.unique().list.len() == 1)
    )
    .explode(["a", "b", "c"])
    .select("a", "b", "c")
)

"""
shape: (2, 3)
┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ str │
╞═════╪═════╪═════╡
│ 4   ┆ 2   ┆ b   │
│ 4   ┆ 2   ┆ f   │
└─────┴─────┴─────┘
"""

Alguém pode sugerir uma maneira melhor de conseguir o mesmo? Lutei um pouco para descobrir essa lógica, então imagino que exista uma maneira mais direta/elegante de obter o mesmo resultado.

python

2 respostas

  1. Best Answer
    • .struct()combinar ae bem uma coluna para que possamos verificar a exclusividade.
    • n_unique()para verificar a exclusividade.
    • over()para limitar o cálculo a ficar dentro de a + b.
    df.filter(
    pl.struct("a","b").n_unique().over(pl.col.a + pl.col.b) == 1
)

┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ str │
╞═════╪═════╪═════╡
│ 4   ┆ 2   ┆ b   │
│ 4   ┆ 2   ┆ f   │
└─────┴─────┴─────┘

    Se você precisar estendê-lo para um número maior de colunas, poderá usá-lo sum_horizontal()para torná-lo mais genérico:

    columns = ["a","b"]

df.filter(
    pl.struct(columns).n_unique().over(pl.sum_horizontal(columns)) == 1
)
    • 4

  2. A partir de

    df = pl.DataFrame(
    {
    "a": [1, 4, 2, 4, 7, 4],
    "b": [4, 2, 3, 3, 0, 2],
    "c": ["a", "b", "c", "d", "e", "f"],
}
)

    Você pode usar an overpara simplificar o group_by e explode, caso contrário, isso não é muito diferente. Ele também usa pl.all_horizontalem vez &de um gerador (para que você possa adicionar facilmente mais colunas).

    (
    df
    .with_columns(ab=pl.col('a')+pl.col('b'))
    .filter(
        pl.all_horizontal((pl.col(x).n_unique()==1).over('ab') for x in ['a','b'])
        )
    .drop('ab')    
)

    Uma observação: isso trataria 2+4 e 4+2 como combinações únicas, mas não sei se isso é pretendido.

    • 1

relate perguntas