如何在Python Polars中获取列表列最大值的索引？含字符串列表转int场景

场景一：列表元素为int类型 #

需求：no_of_docs列是list[int]类型，新增一列存储每个列表中最大值的索引。

示例DataFrame #

import polars as pl

df = pl.DataFrame({
    'no_of_docs': [
        [9,4,2],
        [3,9,1,10],
        [10,3,2,1],
        [10,30],
        [1,2,3,6,4,5]
    ]
})

解决方案 #

使用Polars的arr.arg_max()方法直接获取每个列表的最大值索引：

result = df.with_columns(
    idx=pl.col('no_of_docs').arr.arg_max()
)

# 调整显示配置，完整展示列表内容
pl.Config(
    fmt_table_cell_list_len=10, 
    fmt_str_lengths=80
)

print(result)

场景二：列表元素为str类型（需转int后处理） #

需求：no_of_docs列是list[str]类型，先将列表内元素转为int类型，再获取每个列表中最大值的索引。

示例DataFrame #

import polars as pl

df = pl.DataFrame({
    'no_of_docs': [
        ['9','4','2'],
        ['3','9','1','10'],
        ['10','3','2','1'],
        ['10','30'],
        ['1','2','3','6','4','5']
    ]
})

解决方案 #

先通过arr.eval()结合cast()转换元素类型，再用arr.arg_max()获取索引：

result = df.with_columns(
    idx=pl.col('no_of_docs')
        .arr.eval(pl.element().cast(pl.Int64))
        .arr.arg_max()
)

# 调整显示配置，完整展示列表内容
pl.Config(
    fmt_table_cell_list_len=10, 
    fmt_str_lengths=80
)

print(result)

预期输出 #

shape: (5, 2)
┌────────────────────────────────┬─────┐
│ no_of_docs                     ┆ idx │
│ ---                            ┆ --- │
│ list[str]                      ┆ u32 │
╞════════════════════════════════╪═════╡
│ ["9", "4", "2"]                ┆ 0   │
│ ["3", "9", "1", "10"]          ┆ 3   │
│ ["10", "3", "2", "1"]          ┆ 0   │
│ ["10", "30"]                   ┆ 1   │
│ ["1", "2", "3", "6", "4", "5"] ┆ 3   │
└────────────────────────────────┴─────┘

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