如何用Python的Polars库为连续出现的1添加唯一标签？

要实现给连续出现的1分配唯一整数ID、0保持为0的需求，可以通过标记新组起始行+累加计数的方式完成，具体步骤如下：

完整代码实现 #

import polars as pl

data = {'Switch': [0, 1, 1, 0, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 1, 0, 0, 1]}
df = pl.DataFrame(data)

result = df.with_columns(
    # 标记每个连续1组的起始行
    is_new_group = pl.when((pl.col("Switch") == 1) & (pl.col("Switch").shift(1) != 1)).then(1).otherwise(0),
    # 累加标记生成唯一组号，0对应的行置为0
    "Switch Event Number" = pl.col("is_new_group").cumsum().cast(pl.UInt32).where(pl.col("Switch") == 1, 0)
).drop("is_new_group")

print(result)

代码解释 #

1. 标记新组起始行：

   • 使用shift(1)获取前一行的Switch值，第一行无前置行时返回null，null !=1会被判定为True，因此第一个出现的1会被正确标记为新组起始。
   • 通过when/then/otherwise构造is_new_group列：当前行是1且前一行不是1时标记为1，其余情况为0。

2. 生成唯一组号：

   • 对is_new_group列做累加(cumsum())，每遇到一个新组起始标记，累加值就会递增，自然形成连续1组的唯一ID。
   • 用where方法将Switch为0的行的组号置为0，最后转成UInt32类型匹配期望输出的格式。

3. 清理中间列：

   • 用drop("is_new_group")移除中间生成的标记列，得到最终结果。

执行代码后，输出会完全匹配你给出的期望格式，后续可以直接用Switch Event Number列进行groupby或pivot操作。

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