如何在分组后的Polars DataFrame中替换每组最后两行的value列值为0

嘿，我来帮你搞定这个需求！要实现按symbol分组后，把每组最后两行的value列替换成0，我们可以借助Polars的窗口函数轻松实现，下面给你两种实用的方法：

首先先看一下原始的数据集：

import polars as pl

df = pl.DataFrame(
    {"symbol": [*["A"] * 4, *["B"] * 4], "value": range(8)}
)

对应的输出是：

shape: (8, 2)
┌────────┬───────┐
│ symbol ┆ value │
│ ---    ┆ ---   │
│ str    ┆ i64   │
╞════════╪═══════╡
│ A      ┆ 0     │
│ A      ┆ 1     │
│ A      ┆ 2     │
│ A      ┆ 3     │
│ B      ┆ 4     │
│ B      ┆ 5     │
│ B      ┆ 6     │
│ B      ┆ 7     │
└────────┴───────┘

方法一：利用反向排名定位最后两行 #

我们可以对每组的value列进行倒序排名，这样每组的最后一行排名为1，倒数第二行排名为2，然后通过条件判断把这两行的value替换成0：

result = df.with_columns(
    pl.when(
        pl.col("value").rank(descending=True, method="ordinal").over("symbol") <= 2
    ).then(0).otherwise(pl.col("value")).alias("value")
)

方法二：利用反向累计计数定位最后两行 #

这种方法更直观，通过反向累计计数（每组从最后一行开始数，最后一行是0，倒数第二行是1），判断计数<=1的行就替换成0：

result = df.with_columns(
    pl.when(
        pl.col("value").cumcount(reverse=True).over("symbol") <= 1
    ).then(0).otherwise(pl.col("value")).alias("value")
)

两种方法运行后都会得到你想要的结果：

shape: (8, 2)
┌────────┬───────┐
│ symbol ┆ value │
│ ---    ┆ ---   │
│ str    ┆ i64   │
╞════════╪═══════╡
│ A      ┆ 0     │
│ A      ┆ 1     │
│ A      ┆ 0     │ # 已替换
│ A      ┆ 0     │ # 已替换
│ B      ┆ 4     │
│ B      ┆ 5     │
│ B      ┆ 0     │ # 已替换
│ B      ┆ 0     │ # 已替换
└────────┴───────┘

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