在内部列表上滚动窗口

Question

我有一个这样的框架：

import polars as pl
from polars import Boolean, List, col

src = pl.DataFrame(
    {
        "c1": ["a", "b", "c", "d"],
        "c2": [
            [0, 0],
            [0, 1, 0, 0],
            [1, 0, 1, 1, 1],
            [1, 1, 0],
        ],
    },
    schema_overrides={"c2": List(Boolean)},
)

对于“c2”中的每个内部列表，我试图计算长度为

3

的滑动窗口中的最大总和：

[1, 0, 0, 1, 1, 1]
⇒
[(1, 0, 0), (0, 0, 1), (0, 1, 1),(1, 1, 1)]
⇒
[1, 1, 2, 3]
⇒
3

最有效的方法是什么？

---

我找不到滚动内部列表的方法。 StackOverflow 上的几个答案建议我

• 分解内部列表，
• 滚动分解的系列，
• 以某种方式恢复结构。

...这会导致如下的怪物：

(
    src.with_row_index("outer_index")
    .explode("c2")
    .with_row_index("inner_index")
    .rolling("inner_index", period="3i", offset="0i", closed="left")
    .agg("outer_index", "c1", col("c2").sum())
    .drop("inner_index")
)

首先，这很丑。这是错误的：滚动窗口必须尊重原始内部列表的边框，但在这里它们不这样做。

为了便于阅读，这是我使用

df.to_dicts()

获得的转储

[{'outer_index': [0, 0, 1], 'c1': ['a', 'a', 'b'], 'c2': 0},
 {'outer_index': [0, 1, 1], 'c1': ['a', 'b', 'b'], 'c2': 1},
 {'outer_index': [1, 1, 1], 'c1': ['b', 'b', 'b'], 'c2': 1},
 {'outer_index': [1, 1, 1], 'c1': ['b', 'b', 'b'], 'c2': 1},
 {'outer_index': [1, 1, 2], 'c1': ['b', 'b', 'c'], 'c2': 1},
 {'outer_index': [1, 2, 2], 'c1': ['b', 'c', 'c'], 'c2': 1},
 {'outer_index': [2, 2, 2], 'c1': ['c', 'c', 'c'], 'c2': 2},
 {'outer_index': [2, 2, 2], 'c1': ['c', 'c', 'c'], 'c2': 2},
 {'outer_index': [2, 2, 2], 'c1': ['c', 'c', 'c'], 'c2': 3},
 {'outer_index': [2, 2, 3], 'c1': ['c', 'c', 'd'], 'c2': 3},
 {'outer_index': [2, 3, 3], 'c1': ['c', 'd', 'd'], 'c2': 3},
 {'outer_index': [3, 3, 3], 'c1': ['d', 'd', 'd'], 'c2': 2},
 {'outer_index': [3, 3], 'c1': ['d', 'd'], 'c2': 1},
 {'outer_index': [3], 'c1': ['d'], 'c2': 0}]

Answer 1

解决您的问题：

1. 定义一个函数

   max_sum

   ，它将列表和窗口大小作为输入。它从输入列表中生成给定大小的所有可能的子列表，计算每个子列表的总和，并返回最大总和。
2. 使用列表理解将

   max_sum

   应用于

   c2

   列中的每个列表。

这是实现：

def max_sum(lst, n):
    return max(sum(lst[i:i+n]) for i in range(len(lst)-n+1))
src['c2'] = [max_sum(lst, 3) for lst in src['c2']]

这会将

c2

中的每个列表替换为该列表中长度为 3 的滑动窗口的最大总和。如果包含的元素少于 3 个，

max_sum

将返回列表中所有元素的总和。

Answer 2

你可以稍微调整一下你的表情，这样就可以了。您需要在

group_by

 方法中添加 

DataFrame.rolling() 参数：

(
    src.with_row_index("outer_index")
    .explode("c2")
    .with_row_index("inner_index")
    .rolling("inner_index", period="3i", offset="0i", closed="left", group_by=['outer_index','c1'])
    .agg(col("c2").sum())
    .drop("inner_index")
    .drop("outer_index")
)

┌─────┬───────────┬─────┐
│ c1  ┆ c2        ┆ s   │
│ --- ┆ ---       ┆ --- │
│ str ┆ list[i64] ┆ i64 │
╞═════╪═══════════╪═════╡
│ a   ┆ [0, 0]    ┆ 0   │
│ a   ┆ [0]       ┆ 0   │
│ b   ┆ [0, 1, 0] ┆ 1   │
│ b   ┆ [1, 0, 0] ┆ 1   │
│ b   ┆ [0, 0]    ┆ 0   │
│ …   ┆ …         ┆ …   │
│ c   ┆ [1, 1]    ┆ 2   │
│ c   ┆ [1]       ┆ 1   │
│ d   ┆ [1, 1, 0] ┆ 2   │
│ d   ┆ [1, 0]    ┆ 1   │
│ d   ┆ [0]       ┆ 0   │
└─────┴───────────┴─────┘