如何在Polars中新增列判断两列表列首元素是否相同？

你可以通过提取两个列表列的首元素进行比较，再将布尔结果转换为整数（1/0）来实现需求，以下是两种简洁的实现方式：

方法1：使用列表索引语法 #

直接通过list[0]提取列表的第一个元素，比较后转换为整数类型：

import polars as pl

pl.Config(fmt_table_cell_list_len=6, fmt_str_lengths=100)

df = pl.DataFrame({
    "foo": [[1, 3, 3, 3], [3, 5, 3, 4], [4, 7, 5, 3]], 
    "bar": [[3, 345, 3, 4], [3, 4, 334, 2], [4, 52, 4, 2]]
})

# 新增common_start列
df = df.with_columns(
    common_start=(pl.col("foo").list[0] == pl.col("bar").list[0]).cast(pl.Int64)
)

print(df)

方法2：使用list.get()方法 #

list.get()允许你设置默认值处理空列表的情况，避免索引越界报错：

df = df.with_columns(
    common_start=(
        pl.col("foo").list.get(0, default=-1) == pl.col("bar").list.get(0, default=-1)
    ).cast(pl.Int64)
)

如果你的数据中存在空列表，可以根据实际需求调整default参数的值。

运行结果 #

执行上述代码后，输出的DataFrame与预期一致：

shape: (3, 3)
┌──────────────┬────────────────┬──────────────┐
│ foo          ┆ bar            ┆ common_start │
│ ---          ┆ ---            ┆ ---          │
│ list[i64]    ┆ list[i64]      ┆ i64          │
╞══════════════╪════════════════╪══════════════╡
│ [1, 3, 3, 3] ┆ [3, 345, 3, 4] ┆ 0            │
│ [3, 5, 3, 4] ┆ [3, 4, 334, 2] ┆ 1            │
│ [4, 7, 5, 3] ┆ [4, 52, 4, 2]  ┆ 1            │
└──────────────┴────────────────┴──────────────┘

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