Есть ли какая-нибудь похожая функция idxmax() в py-polars в groupby? ⇐ Python

[Anonymous]

Код: Выделить всё

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()]

кажется, я не могу получить то, что хочу..
есть ли способ сгенерировать столбец «top» в полярах?
Спасибо большое!

Подробнее здесь: https://stackoverflow.com/questions/708 ... in-groupby