Polars列表分组交集方法详解

对于一个文章开发者来说，牢固扎实的基础是十分重要的，golang学习网就来带大家一点点的掌握基础知识点。今天本篇文章带大家了解《Polars列表分组求交集方法与技巧》，主要介绍了，希望对大家的知识积累有所帮助，快点收藏起来吧，否则需要时就找不到了！

本文探讨了在 Polars 中对包含字符串列表的列进行分组求交集的有效方法。针对直接使用列表操作可能遇到的挑战，教程提出了一种通过展开（explode）、行索引计数和过滤的巧妙策略。该方法将列表交集问题转化为元素在组内所有原始行中出现的计数问题，最终实现高效且准确的分组列表交集聚合。

在数据处理中，我们经常会遇到需要对 DataFrame 中的列表（list）类型列进行分组聚合的场景。一个常见的需求是，对于每个分组，找出其所有列表中元素的交集。例如，给定一个包含 id 和 values（字符串列表）的 DataFrame，我们希望得到每个 id 组内所有 values 列表的共同元素。

考虑以下示例 DataFrame：

import polars as pl

df = pl.DataFrame(
   {"id": [1,1,2,2,3,3],
    "values": [["A", "B"], ["B", "C"], ["A", "B"], ["B", "C"], ["A", "B"], ["B", "C"]]
   }
)
print(df)

其输出为：

shape: (6, 2)
┌─────┬───────────┐
│ id  ┆ values    │
│ --- ┆ ---       │
│ i64 ┆ list[str] │
╞═════╪═══════════╡
│ 1   ┆ ["A", "B"]│
│ 1   ┆ ["B", "C"]│
│ 2   ┆ ["A", "B"]│
│ 2   ┆ ["B", "C"]│
│ 3   ┆ ["A", "B"]│
│ 3   ┆ ["B", "C"]│
└─────┴───────────┘

我们期望的输出是每个 id 组内 values 列表的交集，即：

shape: (3, 2)
┌─────┬───────────┐
│ id  ┆ values    │
│ --- ┆ ---       │
│ i64 ┆ list[str] │
╞═════╪═══════════╡
│ 1   ┆ ["B"]     │
│ 2   ┆ ["B"]     │
│ 3   ┆ ["B"]     │
└─────┴───────────┘

直接列表操作的局限性

初看起来，我们可能会尝试使用 pl.reduce 结合 Polars 提供的列表集合操作 list.set_intersection。然而，这种方法往往无法直接达到预期效果，因为它可能在 Series 级别而不是列表元素级别进行聚合，导致结果不符合预期。

例如，以下尝试：

# 尝试1：直接在聚合中使用pl.reduce
result_reduce_1 = df.group_by("id").agg(
    pl.reduce(function=lambda acc, x: acc.list.set_intersection(x),
              exprs=pl.col("values"))
)
print(result_reduce_1)

输出显示 values 列变成了 list[list[str]]，这并非我们想要的单个列表的交集：

shape: (3, 2)
┌─────┬──────────────────────────┐
│ id  ┆ values                   │
│ --- ┆ ---                      │
│ i64 ┆ list[list[str]]          │
╞═════╪══════════════════════════╡
│ 1   ┆ [["A", "B"], ["B", "C"]] │
│ 3   ┆ [["A", "B"], ["B", "C"]] │
│ 2   ┆ [["A", "B"], ["B", "C"]] │
└─────┴──────────────────────────┘

另一个尝试是先 explode 再 reduce，但这会将所有元素扁平化后再进行集合操作，失去了原始列表的边界信息，导致结果是所有元素的并集而非交集：

今天关于《Polars列表分组交集方法详解》的内容介绍就到此结束，如果有什么疑问或者建议，可以在golang学习网公众号下多多回复交流；文中若有不正之处，也希望回复留言以告知！