如何在Polars中逐行应用掩码?

后端开发 2026-07-12

让我们假设要逐行找出连续1 的最大长度,下面的df如下。

import polars as pl

df = pl.from_repr("""
┌─────┬─────┬─────┬─────┬─────┐
│ 0   ┆ 1   ┆ 2   ┆ 3   ┆ 4   │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╪═════╡
│ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 0   │
│ 1   ┆ 0   ┆ 0   ┆ 1   ┆ 1   │
│ 0   ┆ 0   ┆ 1   ┆ 0   ┆ 1   │
│ 1   ┆ 0   ┆ 0   ┆ 0   ┆ 0   │
│ 1   ┆ 0   ┆ 1   ┆ 1   ┆ 1   │
│ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 0   │
│ 1   ┆ 1   ┆ 1   ┆ 0   ┆ 0   │
│ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 0   │
│ 1   ┆ 1   ┆ 0   ┆ 1   ┆ 1   │
│ 1   ┆ 1   ┆ 1   ┆ 0   ┆ 1   │
└─────┴─────┴─────┴─────┴─────┘
""")

在pandas中,我们可以这样做:

m = df.to_pandas().eq(1)
m.cumsum(axis=1).sub(m.cumsum(axis=1).mask(m).ffill(axis=1).fillna(0)).max(axis=1)
0    0.0
1    2.0
2    1.0
3    1.0
4    3.0
5    0.0
6    3.0
7    0.0
8    2.0
9    3.0
dtype: float64

用Polars怎么实现同样的操作?

解决方案

与你在 另一个问题 中展示的方法类似,你可以 concat_arr 每一行。

.arr.agg() 可以用来执行 .rle(),这会得到一个列表,你可以对其进行筛选并 .agg() 最大值。

df.with_columns(
    pl.concat_arr(pl.all() == 1)
      .arr.agg(pl.element().rle())
      .list.filter(pl.element().struct.field("value"))
      .list.agg(pl.element().struct.field("len").max().fill_null(0))
      .alias("rle_max")
)
# shape: (10, 6)
# ┌─────┬─────┬─────┬─────┬─────┬─────────┐
# │ 0   ┆ 1   ┆ 2   ┆ 3   ┆ 4   ┆ rle_max │
# │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ ---     │
# │ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 ┆ u32     │
# ╞═════╪═════╪═════╪═════╪═════╪═════════╡
# │ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 0       │
# │ 1   ┆ 0   ┆ 0   ┆ 1   ┆ 1   ┆ 2       │
# │ 0   ┆ 0   ┆ 1   ┆ 0   ┆ 1   ┆ 1       │
# │ 1   ┆ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 1       │
# │ 1   ┆ 0   ┆ 1   ┆ 1   ┆ 1   ┆ 3       │
# │ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 0       │
# │ 1   ┆ 1   ┆ 1   ┆ 0   ┆ 0   ┆ 3       │
# │ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 0   ┆ 0       │
# │ 1   ┆ 1   ┆ 0   ┆ 1   ┆ 1   ┆ 2       │
# │ 1   ┆ 1   ┆ 1   ┆ 0   ┆ 1   ┆ 3       │
# └─────┴─────┴─────┴─────┴─────┴─────────┘

你可以使用 pl.concat_list().list.agg() 来避免混合数组和列表,但在我的测试中,从 concat_arr 开始要稍微“快”一些。

如果我们用一个Series的示例来展示每一步。

rle 给出了长度

s = pl.Series([1, 1, 1, 0, 1, 1, 1, 1]).cast(bool)

s.rle()

# shape: (3,)
# Series: '' [struct[2]]
# [
#     {3,true}
#     {1,false}
#     {4,true}
# ]

然后从中筛选出true的值

s.rle().filter(s.rle().struct.field("value"))

# shape: (2,)
# Series: '' [struct[2]]
# [
#     {3,true}
#     {4,true}
# ]

并再取最大长度

s.rle().filter(s.rle().struct.field("value")).struct.field("len").max()

# 4
站内所有文章版权归属LeftHeroAI导航站,无授权禁止任何主体转载、抄袭、复制内容,亦不得私自架设镜像站点。一经侵权,本站将通过法律途径追责。

相关文章