如何用更简洁的Polars方式实现：为DataFrame添加列以查找后续首个满足价格≥当前限制的最小索引

嘿，我来帮你把这个Polars的实现简化一下！你的需求是给每行找到之后第一个价格大于等于当前行限制的最小索引，其实不用拆分成8步，用Polars的链式操作就能搞定，这里有两种不错的方案，你可以根据数据量和场景选择：

方案一：简洁直观的交叉连接法（适合中小数据集） #

这个方案用自连接+筛选+分组取最小的思路，代码非常简洁，逻辑也一目了然：

import polars as pl

# 原DataFrame
df_1 = pl.DataFrame({
    'name': ['Alpha', 'Alpha', 'Alpha', 'Alpha', 'Alpha'],
    'index': [0, 3, 4, 7, 9],
    'limit': [12, 18, 11, 5, 9],
    'price': [10, 15, 12, 8, 11]
})

# 链式操作一步到位
df_final = (
    # 按name自连接，生成所有同组行的配对
    df_1.join(df_1, on="name", suffix="_next")
    # 筛选出「后续行index更大」且「后续行price≥当前行limit」的配对
    .filter(pl.col("index_next") > pl.col("index"), pl.col("price_next") >= pl.col("limit"))
    # 按原行的所有字段分组，取符合条件的最小后续index
    .group_by(["name", "index", "limit", "price"])
    .agg(pl.col("index_next").min().alias("min_index"))
    # 右连接回原表，确保所有原行都保留，无符合条件的自动填null
    .join(df_1, on=["name", "index", "limit", "price"], how="right")
)

print(df_final)

运行后输出的结果和你预期的完全一致，整个过程只用了4步链式操作，没有冗余的中间变量。

方案二：高效的组内列表处理法（适合大数据集） #

如果你的数据集很大，交叉连接会生成O(n²)的临时行，可能影响性能。这时候可以用组内列表处理的方式，只在每个name组内部操作，效率更高：

df_final = (
    df_1.group_by("name")
    # 给每行标记组内位置，同时把组内的price和index存成列表
    .with_columns(
        row_pos=pl.int_range(0, pl.count()),
        prices=pl.col("price").list(),
        indices=pl.col("index").list()
    )
    # 对每行，从当前位置之后的元素里找第一个满足条件的index
    .with_columns(
        min_index=pl.struct("row_pos", "limit", "prices", "indices").map_elements(
            lambda x: 
                next(
                    (idx for p, idx in zip(x["prices"][x["row_pos"]+1:], x["indices"][x["row_pos"]+1:]) if p >= x["limit"]),
                    None
                )
        )
    )
    # 去掉临时辅助列
    .drop("row_pos", "prices", "indices")
)

print(df_final)

这个方法通过把组内数据转成列表，只遍历每个组一次，避免了大量的临时行，数据量越大，性能优势越明显。

两种方案都能完美实现你的需求，你可以根据自己的数据集大小来选择～

备注：内容来源于stack exchange，提问作者Danilo Setton