如何在Polars数据框中对每行最大的x列进行行聚合并生成列标识矩阵

这个需求我之前也碰到过，核心就是要把数值和对应的列名绑定起来处理，这样既能算出top3的和，又能标记出哪些列被选中。我来给你一步步拆解怎么实现：

首先，先明确我们要处理的是所有以val开头的列，先把它们单独拎出来方便批量操作。然后通过Polars的struct类型把每个列的数值和列名绑定，这样后续排序取top3的时候，能同时拿到列名信息。

完整实现代码 #

import polars as pl

# 初始化你的原始数据框
df = pl.DataFrame({
        'id_col': [0,1,2,3,4],
        'val1': [10,0,0,20,5],
        'val2': [5,1,2,3,10],
        'val3': [8,2,2,2,5],
        'val4': [1,7,7,4,1],
        'val5': [3,0,0,6,0],
        'val6': [2,7,5,5,4]
})

# 筛选出所有值列（val1到val6）
val_cols = [col for col in df.columns if col.startswith('val')]

# 一步步处理得到结果
result = df.with_columns(
    # 把每个值列转换成「数值+列名」的struct，再合并成列表列
    pl.concat_list(
        [pl.struct(pl.col(c).alias("value"), pl.lit(c).alias("col_name")) for c in val_cols]
    ).alias("struct_list")
).with_columns(
    # 对列表按数值降序排序，取前3个元素（即每行最大的3个值对应的struct）
    pl.col("struct_list").list.sort_by("value", descending=True).list.slice(0, 3).alias("top3")
).with_columns(
    # 计算top3的数值之和，得到agg_col
    pl.col("top3").list.eval(pl.element().struct.field("value")).list.sum().alias("agg_col"),
    # 提取top3对应的列名，存成列表
    pl.col("top3").list.eval(pl.element().struct.field("col_name")).alias("selected_cols")
).with_columns(
    # 对每个值列，判断是否在选中的列名列表里，生成1/0的标识
    [pl.col("selected_cols").list.contains(c).cast(pl.Int8).alias(c) for c in val_cols]
).drop("struct_list", "top3", "selected_cols")  # 删掉中间临时列

print(result)

关键步骤解释 #

1. 绑定数值与列名：用pl.struct把每个值列的数值和对应的列名打包成结构体，这样后续排序后不会丢失列名信息——这解决了你之前只排序数值但没法关联原列的问题。
2. 取top3元素：通过list.sort_by按数值降序排序，再用list.slice(0,3)截取前3个最大的元素。如果有并列的情况，Polars会按原列的顺序保留，完全符合你说的「任意选择并列项」的要求。
3. 生成标识矩阵：用list.contains判断每个值列的名称是否在top3的列名列表里，再转成整数类型（1表示选中，0表示未选中），直接替换原有的值列即可。

运行这段代码后，就能得到你想要的结果：既有标记选中列的矩阵，又有每行top3的求和列agg_col。

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