如何用Polars生成组合（非排列）的DataFrame？

先跟你明确下组合和排列的核心区别：比如我们有一组基础元素{a, b}，如果考虑元素顺序（也就是排列），会有[a,a]、[a,b]、[b,a]、[b,b]这些可能；要是不允许重复元素，排列就是[a,b]和[b,a]；但组合是完全不关心元素顺序的，所以[a,b]和[b,a]会被视为同一个组合，只需要保留一份。

你现在需要生成这种不考虑顺序的组合，目前已经给出了一个可行的Polars实现方案，我们先来看下你的代码逻辑，再聊聊有没有更优化的实现方式：

你的现有实现方案 #

import polars as pl

choices = pl.DataFrame(
    [
        pl.Series("flavor", ["x"] * 2 + ["y"] * 3),
        pl.Series("choice", ["a", "b"] + ["1", "2", "3"]),
    ]
)

# join to produce the choices
choices.join(choices, on=["flavor"]).with_columns(
    # generate a 2-element list representing the choice
    sorted_choice_pair=pl.concat_list("choice", "choice_right").list.sort()
).filter(pl.col.choice.eq(pl.col.sorted_choice_pair.list.first()))

这个方案的思路很清晰：

1. 先按flavor字段进行自连接，生成同口味下所有可能的元素配对（这一步会包含所有排列，比如[a,b]和[b,a]都会被生成）
2. 把每一对choice和choice_right合并成列表并排序，这样不管原配对的顺序如何，相同组合的排序结果都是完全一致的
3. 最后过滤出choice等于排序后列表第一个元素的行，这样就把重复的组合只保留一份，得到我们需要的结果

运行这段代码后，得到的结果完全符合组合的要求：

shape: (9, 4)
┌────────┬────────┬──────────────┬────────────────────┐
│ flavor ┆ choice ┆ choice_right ┆ sorted_choice_pair │
│ ---    ┆ ---    ┆ ---          ┆ ---                │
│ str    ┆ str    ┆ str          ┆ list[str]          │
╞════════╪════════╪══════════════╪════════════════════╡
│ x      ┆ a      ┆ a            ┆ ["a", "a"]         │
│ x      ┆ a      ┆ b            ┆ ["a", "b"]         │
│ x      ┆ b      ┆ b            ┆ ["b", "b"]         │
│ y      ┆ 1      ┆ 1            ┆ ["1", "1"]         │
│ y      ┆ 1      ┆ 2            ┆ ["1", "2"]         │
│ y      ┆ 2      ┆ 2            ┆ ["2", "2"]         │
│ y      ┆ 1      ┆ 3            ┆ ["1", "3"]         │
│ y      ┆ 2      ┆ 3            ┆ ["2", "3"]         │
│ y      ┆ 3      ┆ 3            ┆ ["3", "3"]         │
└────────┴────────┴──────────────┴────────────────────┘

更高效的优化方案 #

你的方案是可行的，但自连接会生成O(n²)的中间结果，当数据量较大时可能会有性能损耗。我们可以换一种思路，在自连接的阶段就直接限制条件，避免生成重复的配对：

import polars as pl

choices = pl.DataFrame(
    [
        pl.Series("flavor", ["x"] * 2 + ["y"] * 3),
        pl.Series("choice", ["a", "b"] + ["1", "2", "3"]),
    ]
)

# 自连接时直接添加过滤条件，只保留choice <= choice_right的配对
result = choices.join(choices, on=["flavor"], how="inner").filter(pl.col("choice") <= pl.col("choice_right"))
print(result)

这个方法的优势很明显：

• 中间结果的行数更少，不需要生成像[b,a]这种和[a,b]重复的配对
• 不需要额外的列表排序操作，减少了计算开销，性能更优
• 结果同样符合组合的要求，因为我们只保留了元素顺序不逆的配对，相当于直接完成了去重

不过要注意，这个方法依赖choice列的元素是可比较的类型（比如字符串、数字等），如果choice是不可比较的复杂类型，那还是需要使用你原来的排序过滤方法。

备注：内容来源于stack exchange，提问作者bzm3r