Início / coding / Perguntas / 77488797
Accepted
asongtoruin
Asked: 2023-11-15 23:00:25 +0800 CST

polars - obtém o nome da coluna contendo o valor máximo por linha

  • 772

Tenho pandasa opção de fazer o seguinte para obter o nome da coluna que contém o valor máximo de cada linha:

import pandas as pd

df = pd.DataFrame({'a': [1, 2, 3, 4, 5], 'b': [5, 4, 3, 2, 1]})

df['Largest'] = df.idxmax(axis=1)

O que me pega:

a b Maior
0 1 5 b
1 2 4 b
2 3 3 a
3 4 2 a
4 5 1 a

Como posso fazer esse tipo de operação em polars? Não parece haver um idxmaxmétodo e max_horizontalparece retornar apenas o valor, e não qualquer informação de indexação.

python-polars

1 respostas

  1. Best Answer

    Você poderia fazer

    def arg_max_horizontal(*columns: pl.Expr) -> pl.Expr:
    return (
        pl.concat_list(columns)
        .list.arg_max()
        .map_dict({i: col_name for i, col_name in enumerate(columns)})
    )


print(dfpl.with_columns(Largest=arg_max_horizontal("a", "b")))

    Então você vai conseguir

    Out[11]:
shape: (5, 3)
┌─────┬─────┬─────────┐
│ a   ┆ b   ┆ Largest │
│ --- ┆ --- ┆ ---     │
│ i64 ┆ i64 ┆ str     │
╞═════╪═════╪═════════╡
│ 1   ┆ 5   ┆ b       │
│ 2   ┆ 4   ┆ b       │
│ 3   ┆ 3   ┆ a       │
│ 4   ┆ 2   ┆ a       │
│ 5   ┆ 1   ┆ a       │
└─────┴─────┴─────────┘
    • 2

relate perguntas