python - groupby 中的 py-polars 中是否有 idxmax() 的类似函数？

Question

import polars as pl
import pandas as pd


A = ['a','a','a','a','a','a','a','b','b','b','b','b','b','b']
B = [1,2,3,4,5,6,7,8,9,10,11,12,13,14]


df = pl.DataFrame({'cola':A,
                   'colb':B})


df_pd = df.to_pandas()

index = df_pd.groupby('cola')['colb'].idxmax()
df_pd.loc[index,'top'] = 1

在 pandas 中，我可以使用 idxmax() 获取顶部的列。

然而，在极地

我使用 arg_max()

index = df[pl.col('colb').arg_max().over('cola').flatten()]

似乎无法得到我想要的..

有什么办法可以在极地中生成一列“顶部”？

多谢！

Answer 1

在 Polars 中，窗口函数 (the .over()) 将执行聚合 + 自连接（参见https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.Expr.over.html? highlight=over#polars.Expr.over），这意味着您不能为每行返回唯一值，这就是您所追求的。

计算顶列的一种方法是使用apply：

df.groupby("cola").apply(lambda x: x.with_columns([pl.col("colb"), (pl.col("colb")==pl.col("colb").max()).alias("top")]))