Python 在Polars中查找给定列的范围内的最大值

我认为一般的 idea 是 for each "范围"分配"组ID".

一种常见的方法是使用累积和和远期填充.

(
    df.with_columns(
        start = (pl.col("Column B") == "Life").cum_sum().forward_fill(),
        end = (pl.col("Column B") == "Death").cum_sum().forward_fill()
    )
    .with_columns(
        group_id_1 = pl.col("start") + pl.col("end")
    )
    .with_columns(
        group_id_2 = 
            pl.when(pl.col("Column B") == "Death")
              .then(None)
              .otherwise(pl.col("group_id_1"))
              .forward_fill()
    )
)

shape: (10, 6)
┌──────────┬──────────┬───────┬─────┬────────────┬────────────┐
│ Column A ┆ Column B ┆ start ┆ end ┆ group_id_1 ┆ group_id_2 │
│ ---      ┆ ---      ┆ ---   ┆ --- ┆ ---        ┆ ---        │
│ i64      ┆ str      ┆ u32   ┆ u32 ┆ u32        ┆ u32        │
╞══════════╪══════════╪═══════╪═════╪════════════╪════════════╡
│ 2        ┆ Life     ┆ 1     ┆ 0   ┆ 1          ┆ 1          │
│ 3        ┆ null     ┆ 1     ┆ 0   ┆ 1          ┆ 1          │
│ 1        ┆ null     ┆ 1     ┆ 0   ┆ 1          ┆ 1          │
│ 4        ┆ null     ┆ 1     ┆ 0   ┆ 1          ┆ 1          │
│ 1        ┆ Death    ┆ 1     ┆ 1   ┆ 2          ┆ 1          │ # 2 -> 1
│ 3        ┆ null     ┆ 1     ┆ 1   ┆ 2          ┆ 2          │
│ 3        ┆ Life     ┆ 2     ┆ 1   ┆ 3          ┆ 3          │
│ 2        ┆ null     ┆ 2     ┆ 1   ┆ 3          ┆ 3          │
│ 1        ┆ null     ┆ 2     ┆ 1   ┆ 3          ┆ 3          │
│ 0        ┆ Death    ┆ 2     ┆ 2   ┆ 4          ┆ 3          │ # 4 -> 3
└──────────┴──────────┴───────┴─────┴────────────┴────────────┘

除了死刑犯之外，group_id_1得到了大部分的结果，死刑犯需要修正才能得到group_id_2

由于它足够复杂，您可能希望使用变量和/或函数来构建最终表达式:

start = pl.col("Column B") == "Life"
end = pl.col("Column B") == "Death"

group_id = (
   pl.when(end)
     .then(None)
     .otherwise((start.cum_sum().forward_fill() + end.cum_sum()).forward_fill())
     .forward_fill()
)

df.with_columns(
   pl.when(start) 
     .then(pl.col("Column A").max())
     .over(group_id)
     .alias("Column C")
)

shape: (10, 3)
┌──────────┬──────────┬──────────┐
│ Column A ┆ Column B ┆ Column C │
│ ---      ┆ ---      ┆ ---      │
│ i64      ┆ str      ┆ i64      │
╞══════════╪══════════╪══════════╡
│ 2        ┆ Life     ┆ 4        │
│ 3        ┆ null     ┆ null     │
│ 1        ┆ null     ┆ null     │
│ 4        ┆ null     ┆ null     │
│ 1        ┆ Death    ┆ null     │
│ 3        ┆ null     ┆ null     │
│ 3        ┆ Life     ┆ 3        │
│ 2        ┆ null     ┆ null     │
│ 1        ┆ null     ┆ null     │
│ 0        ┆ Death    ┆ null     │
└──────────┴──────────┴──────────┘