如何在Polars中逐行应用掩码?
让我们假设要逐行找出连续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导航站,无授权禁止任何主体转载、抄袭、复制内容,亦不得私自架设镜像站点。一经侵权,本站将通过法律途径追责。