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

Question

我有以下数据框：

\n

df = pl.DataFrame({\n    "Column A": [2, 3, 1, 4, 1, 3, 3, 2, 1, 0],\n    "Column B": [\n        "Life", None, None, None, "Death", None, \n        "Life", None, None, "Death"\n    ]\n})\n

\n

shape: (10, 2)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 Column A \xe2\x94\x86 Column B \xe2\x94\x82\n\xe2\x94\x82 ---      \xe2\x94\x86 ---      \xe2\x94\x82\n\xe2\x94\x82 i64      \xe2\x94\x86 str      \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 2        \xe2\x94\x86 Life     \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 4        \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 Death    \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 Life     \xe2\x94\x82\n\xe2\x94\x82 2        \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 0        \xe2\x94\x86 Death    \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

我想创建一个新列，我们将其称为 C 列。对于 B 列为“Life”的每一行，C 列应该具有从该行到 A 列中的值范围内的最大值B 列是“死亡”的行。如果 B 列不是“Life”，则 C 列应设置为“None”

\n

最终结果应该是这样的：

\n

shape: (10, 3)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 Column A \xe2\x94\x86 Column B \xe2\x94\x86 Column C \xe2\x94\x82\n\xe2\x94\x82 ---      \xe2\x94\x86 ---      \xe2\x94\x86 ---      \xe2\x94\x82\n\xe2\x94\x82 i64      \xe2\x94\x86 str      \xe2\x94\x86 f64      \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 2        \xe2\x94\x86 Life     \xe2\x94\x86 4.0      \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 4        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 Death    \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 Life     \xe2\x94\x86 3.0      \xe2\x94\x82\n\xe2\x94\x82 2        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 0        \xe2\x94\x86 Death    \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

如何在 Python 中使用 Polars 来实现这一目标？任何帮助或建议将不胜感激！

\n

Answer 1

我认为总体思路是将“组 ID”分配给每个“范围”。

\n

一种常见的方法是使用累积和以及前向填充。

\n

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

\n

shape: (10, 6)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 Column A \xe2\x94\x86 Column B \xe2\x94\x86 start \xe2\x94\x86 end \xe2\x94\x86 group_id_1 \xe2\x94\x86 group_id_2 \xe2\x94\x82\n\xe2\x94\x82 ---      \xe2\x94\x86 ---      \xe2\x94\x86 ---   \xe2\x94\x86 --- \xe2\x94\x86 ---        \xe2\x94\x86 ---        \xe2\x94\x82\n\xe2\x94\x82 i64      \xe2\x94\x86 str      \xe2\x94\x86 u32   \xe2\x94\x86 u32 \xe2\x94\x86 u32        \xe2\x94\x86 u32        \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 2        \xe2\x94\x86 Life     \xe2\x94\x86 1     \xe2\x94\x86 0   \xe2\x94\x86 1          \xe2\x94\x86 1          \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 null     \xe2\x94\x86 1     \xe2\x94\x86 0   \xe2\x94\x86 1          \xe2\x94\x86 1          \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 null     \xe2\x94\x86 1     \xe2\x94\x86 0   \xe2\x94\x86 1          \xe2\x94\x86 1          \xe2\x94\x82\n\xe2\x94\x82 4        \xe2\x94\x86 null     \xe2\x94\x86 1     \xe2\x94\x86 0   \xe2\x94\x86 1          \xe2\x94\x86 1          \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 Death    \xe2\x94\x86 1     \xe2\x94\x86 1   \xe2\x94\x86 2          \xe2\x94\x86 1          \xe2\x94\x82 # 2 -> 1\n\xe2\x94\x82 3        \xe2\x94\x86 null     \xe2\x94\x86 1     \xe2\x94\x86 1   \xe2\x94\x86 2          \xe2\x94\x86 2          \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 Life     \xe2\x94\x86 2     \xe2\x94\x86 1   \xe2\x94\x86 3          \xe2\x94\x86 3          \xe2\x94\x82\n\xe2\x94\x82 2        \xe2\x94\x86 null     \xe2\x94\x86 2     \xe2\x94\x86 1   \xe2\x94\x86 3          \xe2\x94\x86 3          \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 null     \xe2\x94\x86 2     \xe2\x94\x86 1   \xe2\x94\x86 3          \xe2\x94\x86 3          \xe2\x94\x82\n\xe2\x94\x82 0        \xe2\x94\x86 Death    \xe2\x94\x86 2     \xe2\x94\x86 2   \xe2\x94\x86 4          \xe2\x94\x86 3          \xe2\x94\x82 # 4 -> 3\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

group_id_1除了需要转移以生产的死囚牢房外，大部分时间都在那里group_id_2

\n

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

\n

start = pl.col("Column B") == "Life"\nend   = pl.col("Column B") == "Death"\n\ngroup_id = (start.cum_sum() + end.cum_sum()).forward_fill() # id_1\ngroup_id = (                                                # id_2\n   pl.when(end)\n     .then(group_id.shift())\n     .otherwise(group_id)\n)\n\n# Insert the max over each group into each Life row\ndf.with_columns(\n   pl.when(start) \n     .then(pl.col("Column A").max().over(group_id))\n     .alias("Column C")\n)\n

\n

shape: (10, 3)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 Column A \xe2\x94\x86 Column B \xe2\x94\x86 Column C \xe2\x94\x82\n\xe2\x94\x82 ---      \xe2\x94\x86 ---      \xe2\x94\x86 ---      \xe2\x94\x82\n\xe2\x94\x82 i64      \xe2\x94\x86 str      \xe2\x94\x86 i64      \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 2        \xe2\x94\x86 Life     \xe2\x94\x86 4        \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 4        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 Death    \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 3        \xe2\x94\x86 Life     \xe2\x94\x86 3        \xe2\x94\x82\n\xe2\x94\x82 2        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 1        \xe2\x94\x86 null     \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x82 0        \xe2\x94\x86 Death    \xe2\x94\x86 null     \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n