如何在DataFrame中获取另一列指定元素在列表列中的索引

解决方案 #

利用Polars的数组方法arr.arg_where结合标量提取，可高效实现需求：

import polars as pl

df = pl.DataFrame({'a': [[1, 2, 3], [4, 5, 2], [6, 2, 7]], 'b': [3, 4, 2]})

# 添加列c，存储b在a中的索引
df = df.with_columns(
    c=pl.col("a").arr.arg_where(pl.element() == pl.col("b")).first()
)

print(df)

代码解释 #

• pl.col("a").arr.arg_where(pl.element() == pl.col("b"))：遍历每行的列表列a，筛选出与对应行b值相等的元素的索引，由于每行a元素唯一，结果为单元素数组。
• .first()：将单元素数组转换为标量，得到目标索引值。

运行结果 #

shape: (3, 3)
┌───────────┬─────┬─────┐
│ a         ┆ b   ┆ c   │
│ ---       ┆ --- ┆ --- │
│ list[i64] ┆ i64 ┆ i64 │
╞═══════════╪═════╪═════╡
│ [1, 2, 3] ┆ 3   ┆ 2   │
│ [4, 5, 2] ┆ 4   ┆ 0   │
│ [6, 2, 7] ┆ 2   ┆ 1   │
└───────────┴─────┴─────┘

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