如何在Polars中基于a列位置将c列值映射为b列对应值？

需求说明 #

给定包含三列列表的Polars DataFrame：

• a：list[f32]类型，每行无重复元素
• b：list[f64]类型，与a的列表长度一一对应
• c：list[f32]类型

需要将c列的每个元素替换为：若元素存在于对应行的a列表中，则取a对应位置的b元素；否则替换为0。

问题分析 #

你尝试的replace函数不适用此场景，因为它要求全局的old值无重复，而这里是每行独立的映射关系，需要针对每行构建a→b的映射逻辑。

解决方案 #

方法1：向量化explode+join+聚合（推荐大数据集）

通过展开列表、建立映射、匹配聚合的方式实现完全向量化操作，性能最优：

import polars as pl

# 可复现数据集
df = pl.DataFrame([
    pl.Series('a', [[1.0, 0.0], [5.0, 6.0, 1.0, 0.0, 3.0, 2.0, 4.0], [1.0, 0.0], [0.0, 1.5, 3.0, 1.0, 2.0, 0.5, 2.5], [5.0, 6.0, 0.0, 3.0, 2.0, 4.0, 1.0]]),
    pl.Series('b', [[0.0011431837667905116, 0.9988568162332095], [0.0002857959416976279, 0.00014289797084881395, 0.2842240640182909, 0.5985995998856817, 0.019291226064589884, 0.09388396684767077, 0.003572449271220349], [0.005287224921406116, 0.9947127750785939], [0.8436696198913975, 0.0038582452129179764, 0.00014289797084881395, 0.10703058016576164, 0.007859388396684767, 0.03701057444984281, 0.00042869391254644185], [0.0012860817376393256, 0.0002857959416976279, 0.4645613032294941, 0.038153758216633325, 0.13561017433552444, 0.007430694484138326, 0.3526721920548728]]),
    pl.Series('c', [[0.0, 0.5, 1.0, 1.5], [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5], [0.0, 0.5, 1.0, 1.5], [0.0, 0.5, 0.5, 1.0, 1.0, 1.5, 1.5, 2.0, 2.0, 2.5, 2.5, 3.0, 3.0, 3.5], [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5]]),
])

result = (
    df
    # 添加行索引，确保聚合后对应原行
    .with_row_index("idx")
    # 展开a和b，建立每行的a→b映射关系
    .explode(["a", "b"])
    # 与展开后的c列左连接，匹配映射
    .join(
        df.with_row_index("idx").explode("c"),
        on=["idx", "a"],
        how="right"
    )
    # 未匹配到的元素填充为0
    .with_columns(pl.col("b").fill_null(0.0))
    # 按行索引聚合，重新组合成列表
    .group_by("idx", maintain_order=True)
    .agg(pl.col("b").alias("new_c"))
    # 合并回原DataFrame
    .join(df.with_row_index("idx"), on="idx", how="right")
    # 调整列顺序，移除临时索引
    .select(pl.exclude("idx"), "new_c")
)

print(result.head())

步骤说明：

1. 添加行索引idx，保证后续聚合能对应到原行
2. 展开a和b列，得到每行中a元素与b元素的一一映射
3. 展开c列并与映射表左连接，匹配每个c元素对应的b值
4. 用fill_null(0.0)替换未匹配到的元素
5. 按idx聚合，将b值重新组合为列表，得到替换后的new_c列

方法2：行级字典映射（适合小数据集）

代码更简洁，通过行级映射实现，适合快速验证逻辑：

result = df.with_columns(
    pl.struct("a", "b", "c").map_elements(
        lambda x: [dict(zip(x["a"], x["b"])).get(val, 0.0) for val in x["c"]],
        return_dtype=pl.List(pl.Float64)
    ).alias("new_c")
)

print(result.head())

步骤说明：

1. 对每行的(a,b,c)结构体，用zip构建a→b的字典映射
2. 遍历c的每个元素，用get方法获取对应值，不存在则返回0.0

结果验证 #

第一行的new_c列结果为[0.9988568162332095, 0.0, 0.0011431837667905116, 0.0]，完全符合需求：

• c列的0.0匹配a列的第二个元素，对应b列的0.998857
• 0.5和1.5未在a列中，替换为0
• 1.0匹配a列的第一个元素，对应b列的0.001143

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