Polars：按分组获取另一列最大值对应的目标列值

嗨，我明白你想要实现的效果了——按A分组，找到每个组里x最大的那一行对应的B值，然后把这个值加到每一行的y列里对吧？咱们用Polars的窗口函数就能轻松搞定，给你两种可行的方案：

方案一：通过排序取首值 #

这种方法的思路是在每个分组内，把B列按照x的降序排列，这样x最大对应的B值就会排在第一位，直接取首值即可：

import polars as pl

# 初始化你的DataFrame
df = pl.DataFrame({'A': ['a0', 'a0', 'a1', 'a1'],
                   'B': ['b1', 'b2', 'b1', 'b2'],
                   'x': [0, 10, 5, 1]})

# 添加y列
result = df.with_columns(
    y=pl.col("B").sort_by("x", descending=True).first().over("A")
)

print(result)

执行后得到的结果完全符合你的预期：

shape: (4, 4)
┌─────┬─────┬─────┬─────┐
│ A   ┆ B   ┆ x   ┆ y   │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 ┆ str │
╞═════╪═════╪═════╪═════╡
│ a0  ┆ b1  ┆ 0   ┆ b2  │
│ a0  ┆ b2  ┆ 10  ┆ b2  │
│ a1  ┆ b1  ┆ 5   ┆ b1  │
│ a1  ┆ b2  ┆ 1   ┆ b1  │
└─────┴─────┴─────┴─────┘

方案二：使用arg_max直接定位索引 #

Polars的arg_max函数可以直接返回分组内x最大值所在的行索引，再通过take取出对应的B值，这种方法更直接：

result = df.with_columns(
    y=pl.col("B").take(pl.col("x").arg_max()).over("A")
)

这个逻辑也很清晰：

• pl.col("x").arg_max()：在每个A分组里，找到x值最大的那一行的位置索引
• pl.col("B").take(...)：根据这个索引取出对应的B值
• over("A")：把这个值广播到当前分组的所有行中

两种方法都能达到你想要的效果，你可以根据自己的习惯选择~

备注：内容来源于stack exchange，提问作者rindis