如何获取Polars数据框列表列中含最多分号元素的索引

问题描述 #

给定如下Polars数据框：

import polars as pl

df = pl.DataFrame({'doc_id':[
    ['83;45;32;65;13','7;8;9'],
    ['9;4;5','4;2;7;3;5;8;10;11'],
    ['1000;2000','76;34;100001;7474;2924'],
    ['100;200','200;100'],
    ['3;4;6;7;10;11','1;2;3;4;5']
]})

doc_id列的每个元素是字符串列表，每个字符串用分号分隔文档ID。需要为每行创建新列len_idx_at，填充该行列表中分号数量最多的元素的索引；若多个元素分号数量相同，优先选择索引0。

解决方案 #

利用Polars的列表表达式和字符串计数功能，可以高效实现需求：

df = df.with_columns(
    len_idx_at=pl.col('doc_id')
    .list.eval(pl.element().str.count(';'))  # 计算每个列表元素的分号数量
    .list.arg_max()  # 获取第一个最大值对应的索引
)

代码解释 #

1. list.eval(pl.element().str.count(';'))：遍历doc_id列的每个列表，对列表内的每个字符串元素，统计其中分号;的数量，生成一个整数列表。例如第一行的["83;45;32;65;13", "7;8;9"]会被转换为[4, 2]。
2. list.arg_max()：对生成的整数列表，返回第一个出现最大值的元素索引。当多个元素分号数量相同时，该方法会默认返回索引最小的那个（即优先选索引0），正好符合需求。

验证结果 #

执行代码后，数据框输出如下：

shape: (5, 2)
┌─────────────────────────────────┬────────────┐
│ doc_id                          ┆ len_idx_at │
│ ---                             ┆ ---        │
│ list[str]                       ┆ i32        │
╞═════════════════════════════════╪════════════╡
│ ["83;45;32;65;13", "7;8;9"]     ┆ 0          │
│ ["9;4;5", "4;2;7;3;5;8;10;11"]  ┆ 1          │
│ ["1000;2000", "76;34;100001;74… ┆ 1          │
│ ["100;200", "200;100"]          ┆ 0          │
│ ["3;4;6;7;10;11", "1;2;3;4;5"]  ┆ 0          │
└─────────────────────────────────┴────────────┘

完全符合预期结果。

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