Início / coding / Perguntas / 78979548
Accepted
Andi
Asked: 2024-09-13 02:45:17 +0800 CST

Criar lista de colunas a partir de nomes de colunas

  • 772

Tenho um simples pl.DataFramecom várias colunas que contêm apenas valores booleanos.

import polars as pl

df = pl.DataFrame(
    {"s1": [True, True, False], "s2": [False, True, True], "s3": [False, False, False]}
)

shape: (3, 3)
┌───────┬───────┬───────┐
│ s1    ┆ s2    ┆ s3    │
│ ---   ┆ ---   ┆ ---   │
│ bool  ┆ bool  ┆ bool  │
╞═══════╪═══════╪═══════╡
│ true  ┆ false ┆ false │
│ true  ┆ true  ┆ false │
│ false ┆ true  ┆ false │
└───────┴───────┴───────┘

Preciso adicionar outra coluna que contenha listas de comprimento variável. Uma lista em qualquer linha individual deve conter o nome da coluna onde os valores das colunas S1, s2, e s3têm um Truevalor.

Aqui está o que realmente estou procurando:

shape: (3, 4)
┌───────┬───────┬───────┬──────────────┐
│ s1    ┆ s2    ┆ s3    │ list         │
│ ---   ┆ ---   ┆ ---   │ ---          │
│ bool  ┆ bool  ┆ bool  │ list[str]    │
╞═══════╪═══════╪═══════╡══════════════╡
│ true  ┆ false ┆ false │ ["s1"]       │
│ true  ┆ true  ┆ false │ ["s1", "s2"] │
│ false ┆ true  ┆ false │ ["s2"]       │
└───────┴───────┴───────┴──────────────┘
python

3 respostas

  1. Best Answer

    Você pode criar uma lista de expressões quando/então e então remover os nulos.

    df.with_columns(
   pl.concat_list(
      pl.when(col).then(pl.lit(col)) for col in df.columns
   )
   .list.drop_nulls()
   .alias("list")
)

    shape: (3, 4)
┌───────┬───────┬───────┬──────────────┐
│ s1    ┆ s2    ┆ s3    ┆ list         │
│ ---   ┆ ---   ┆ ---   ┆ ---          │
│ bool  ┆ bool  ┆ bool  ┆ list[str]    │
╞═══════╪═══════╪═══════╪══════════════╡
│ true  ┆ false ┆ false ┆ ["s1"]       │
│ true  ┆ true  ┆ false ┆ ["s1", "s2"] │
│ false ┆ true  ┆ false ┆ ["s2"]       │
└───────┴───────┴───────┴──────────────┘

    Ou.list.eval()

    df.with_columns(
   pl.concat_list(pl.all())
     .list.eval(pl.lit(df.columns).gather(pl.element().arg_true()))
     .alias("list")
)

    shape: (3, 4)
┌───────┬───────┬───────┬──────────────┐
│ s1    ┆ s2    ┆ s3    ┆ list         │
│ ---   ┆ ---   ┆ ---   ┆ ---          │
│ bool  ┆ bool  ┆ bool  ┆ list[str]    │
╞═══════╪═══════╪═══════╪══════════════╡
│ true  ┆ false ┆ false ┆ ["s1"]       │
│ true  ┆ true  ┆ false ┆ ["s1", "s2"] │
│ false ┆ true  ┆ false ┆ ["s2"]       │
└───────┴───────┴───────┴──────────────┘
    • 1

  2. Primeiro, você pode criar struct de todas as colunas usando struct, então você pode usar map_elementspara aplicar uma função lambda para cada linha. A função lambda cria uma compreensão de lista que inclui o nome da coluna colse seu valor correspondente valfor True. Então use alias para a nova coluna como list.

    desired_output = df.with_columns(
    [
        pl.struct(df.columns)
        .map_elements(lambda x: [col for col, val in x.items() if val])
        .alias("list")
    ]
)

    Alternativamente, você pode fazer isso

    result = (
    df.with_columns(
        [
            pl.concat_list(
                [
                    pl.when(pl.col(col)).then(pl.lit(col)).otherwise(pl.lit(None))
                    for col in df.columns
                ]
            ).alias("temp_list")
        ]
    )
    .with_columns(
        [
            pl.col("temp_list")
            .list.eval(pl.element().filter(pl.element().is_not_null()))
            .alias("list")
        ]
    )
    .drop("temp_list")
)

    imprimir(saída_desejada)

    shape: (3, 4)
┌───────┬───────┬───────┬──────────────┐
│ s1    ┆ s2    ┆ s3    ┆ list         │
│ ---   ┆ ---   ┆ ---   ┆ ---          │
│ bool  ┆ bool  ┆ bool  ┆ list[str]    │
╞═══════╪═══════╪═══════╪══════════════╡
│ true  ┆ false ┆ false ┆ ["s1"]       │
│ true  ┆ true  ┆ false ┆ ["s1", "s2"] │
│ false ┆ true  ┆ false ┆ ["s2"]       │
└───────┴───────┴───────┴──────────────┘
    • 0

  3. Você poderia melt, filter, group_by.agg, e join:

    (df.with_row_index()
   .join(df.with_row_index().melt('index').filter(pl.col('value'))
           .group_by('index').agg(pl.col('variable')),
         on='index', how='left'
        )
)

    Saída:

    ┌───────┬───────┬───────┬───────┬──────────────┐
│ index ┆ s1    ┆ s2    ┆ s3    ┆ variable     │
│ ---   ┆ ---   ┆ ---   ┆ ---   ┆ ---          │
│ u32   ┆ bool  ┆ bool  ┆ bool  ┆ list[str]    │
╞═══════╪═══════╪═══════╪═══════╪══════════════╡
│ 0     ┆ true  ┆ false ┆ false ┆ ["s1"]       │
│ 1     ┆ true  ┆ true  ┆ false ┆ ["s1", "s2"] │
│ 2     ┆ false ┆ true  ┆ false ┆ ["s2"]       │
└───────┴───────┴───────┴───────┴──────────────┘
    • 0

relate perguntas