Polars中高效查找col_b首次大于col_a的索引（替代慢循环）

需求说明 #

需要找到每个位置起，col_b的值首次大于col_a值的索引。当前使用for循环实现速度极慢，希望通过窗口函数优化。

示例数据 #

df = polars.DataFrame({"idx": [i for i in range(5)], "col_a": [1,2,3,4,4], "col_b": [1,1,5,5,3]})

# 期望结果：每个位置对应的首个符合条件的索引
result = polars.Series([2,2,2,3,None])

优化思路（伪代码） #

• 按窗口长度比较两列值，使用arg_min()获取索引
• 若未找到结果（如值为None或0），增大窗口长度进行二次扫描
• 重复扫描直到达到最大窗口长度

当前for循环实现 #

df = polars.DataFrame({"col_a": [1,2,3,4,4], "col_b": [1,1,5,5,3]})

for i in range(0, df.shape[0]):
    # `arg_max()`在未找到索引或实际索引为0时均返回0
    series = (df.select("col_a")[i,0] < df.select("col_b")[i:])[:,0]
    idx_found = True in series
    
    if idx_found:
        print(i + series.arg_max())
    else:
        print("None")

# 输出：
2
2
2
3
None

初步优化尝试（编辑1） #

以下代码接近解决问题，但无法区分arg_max()返回的0是实际索引还是未找到结果。后续计划对返回结果为0的位置，使用更长窗口进行二次扫描：

df_res = df.groupby_dynamic("idx", every="1i", period="5i").agg(
    [
       (polars.col("col_a").head(1) < polars.col("col_b")).arg_max().alias("res")
    ]
)

最终解决方案（编辑2） #

先执行编辑1中的代码生成初始结果，再对结果为0的位置使用更大窗口进行后续扫描：

increase_window_size = "10i"
df_res = df_res.groupby_dynamic("idx", every="1i", period=increase_window_size).agg(
    [
       (polars.col("col_a").head(1) < polars.col("col_b")).filter(polars.col("res").head(1) == 0).arg_max().alias("res")
    ]
)

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