如何在Polars中实现滚动窗口内指定值的计数？

问题场景 #

Polars提供rolling_mean()、rolling_apply()、rolling_max()等滚动函数，但直接用rolling_sum()处理布尔列会报错。我们需要统计滚动窗口内特定值（如布尔值True）的出现次数，给定如下LazyFrame：

df = pl.LazyFrame({"Date": ["2023-01-01", "2023-01-02", "2023-01-03", "2023-01-05", "2023-01-10", "2023-01-11", "2023-01-12"], "Pattern": [True, True, False, True, False, False, True]})

要求窗口大小n=3，统计True的出现次数，对应Pandas的实现代码为：

df["Pattern"].apply(lambda x: x == True).rolling(3, min_periods = 0).sum()

Polars的核心解决方案 #

Polars中布尔列无法直接用于rolling_sum()，只需先将布尔类型转为整数（True映射为1，False映射为0），再调用滚动求和函数即可实现计数：

result = df.with_columns(
    pl.col("Pattern").cast(pl.Int32).rolling_sum(window_size=3, min_periods=0).alias("True_Count")
).collect()

执行后得到的True_Count列即为每个窗口内True的出现次数，结果与Pandas一致。

扩展到统计False的次数 #

有两种常用方式：

• 方式一：取反后转整数求和
  将布尔列取反，再转为整数后滚动求和，直接得到False的次数：

  result = df.with_columns(
      pl.col("Pattern").not_().cast(pl.Int32).rolling_sum(window_size=3, min_periods=0).alias("False_Count")
  ).collect()
  

• 方式二：总计数减True计数
  先统计窗口内总元素数和True的次数，再做差值计算：

  result = df.with_columns(
      pl.col("Pattern").cast(pl.Int32).rolling_sum(window_size=3, min_periods=0).alias("True_Count"),
      pl.col("Pattern").rolling_count(window_size=3, min_periods=0).alias("Total_Count")
  ).with_columns(
      (pl.col("Total_Count") - pl.col("True_Count")).alias("False_Count")
  ).collect()
  

扩展到分类数据统计 #

对于分类类型的列，只需先创建一个标识列（判断每行是否等于目标分类值），再转整数后滚动求和即可。

单分类值统计示例 #

假设我们有分类列Category，统计窗口内"A"的出现次数：

df = pl.LazyFrame({"Date": ["2023-01-01", "2023-01-02", "2023-01-03", "2023-01-05", "2023-01-10", "2023-01-11", "2023-01-12"], "Category": ["A", "B", "A", "C", "B", "B", "A"]})

result = df.with_columns(
    (pl.col("Category") == "A").cast(pl.Int32).rolling_sum(window_size=3, min_periods=0).alias("A_Count")
).collect()

多分类值批量统计示例 #

若要同时统计多个分类值的出现次数，可通过列表推导式批量生成计算列：

target_categories = ["A", "B", "C"]
result = df.with_columns(
    [(pl.col("Category") == cat).cast(pl.Int32).rolling_sum(window_size=3, min_periods=0).alias(f"{cat}_Count") for cat in target_categories]
).collect()

内容的提问来源于stack exchange，提问作者SpeedowaGONE