Início / coding / Perguntas / 76969939
Accepted
craigm
Asked: 2023-08-24 21:45:00 +0800 CST

Crie uma nova coluna baseada parcialmente em outros nomes de colunas

  • 772

Sou novo em Polars e Python em geral. Tenho um problema um tanto incomum para o qual preciso de ajuda. Eu tenho um dataframe com mais de 50 colunas que são 0/1. Preciso criar uma nova coluna que contenha uma lista separada por vírgulas de cada coluna que contenha 1, mas usando parte do nome da coluna. Se hccx = 1, anexe x a uma coluna de string. Um exemplo simplificado:

df=pl.DataFrame(
    {'id':[1,2,3], 'hcc1':[0,1,1],'hcc2':[0,0,1],'hcc5':[0,1,1],'hcc8':[1,0,0]}
)

shape: (3, 5)
┌─────┬──────┬──────┬──────┬──────┐
│ id  ┆ hcc1 ┆ hcc2 ┆ hcc5 ┆ hcc8 │
│ --- ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ i64 ┆ i64  ┆ i64  ┆ i64  ┆ i64  │
╞═════╪══════╪══════╪══════╪══════╡
│ 1   ┆ 0    ┆ 0    ┆ 0    ┆ 1    │
│ 2   ┆ 1    ┆ 0    ┆ 1    ┆ 0    │
│ 3   ┆ 1    ┆ 1    ┆ 1    ┆ 0    │
└─────┴──────┴──────┴──────┴──────┘

Quero criar uma nova coluna (tipo string), hccall, semelhante a esta:

eu ia ligar
1 8
2 1,5
3 1,2,5

Imagino que algum tipo de compreensão de lista percorrendo colunas que começam com 'hcc' funcionaria, mas estou meio preso. Posso criar um loop, mas não sei como anexar à coluna dentro do loop. Alguma ideia inteligente?

python-polars

2 respostas

  1. Best Answer

    Acho que a opção mais fácil é primeiro derreter seu dataframe para que você tenha uma linha por ID por coluna em seu dataframe. Você pode filtrar as linhas iguais a 1. Você pode agregar a última letra das linhas originais em uma lista. Com list.join você pode combinar a lista em uma string separada por vírgulas

    (
    df.melt(id_vars="id")
    .filter(pl.col("value") == 1)
    .groupby("id")
    .agg(pl.col("variable").str.slice(3))
    .with_columns(pl.col("variable").list.join(","))
)

shape: (3, 2)
┌─────┬──────────┐
│ id  ┆ variable │
│ --- ┆ ---      │
│ i64 ┆ str      │
╞═════╪══════════╡
│ 2   ┆ 1,5      │
│ 3   ┆ 1,2,5    │
│ 1   ┆ 8        │
└─────┴──────────┘
    • 3

  2. Aqui está outra maneira de fazer isso sem derreter ou agrupar

    (
    df
    .with_columns(
        hccall=pl.concat_list(
                    pl.when(pl.col(x)==1)
                    .then(pl.lit(x.replace('hcc',', ')))
                    .otherwise(pl.lit('')) for x in df.columns if x[:3]=='hcc')
                .list.join("")
                .str.replace("^, ", "")
                    )
)

    Não consigo descobrir como eliminar elegantemente os nulos, concat_listentão fica complicado str.replaceno final. Acho que há uma maneira sucinta de fazer isso, mas isso me escapa no momento.

    • 2

relate perguntas