如何按条件筛选满足a+b对应唯一a、b组合的DataFrame行

数据集 #

data = {
    "a": [1, 4, 2, 4, 7, 4],
    "b": [4, 2, 3, 3, 0, 2],
    "c": ["a", "b", "c", "d", "e", "f"],
}

需求 #

保留那些**a + b的和对应唯一a与b组合**的行——即某个a+b的和只能匹配一组(a,b)值，所有属于该组合的行都需要保留。

现有实现代码 #

df = (
    pl.DataFrame(data)
    .with_columns(sum_ab=pl.col("a") + pl.col("b"))
    .group_by("sum_ab")
    .agg(pl.col("a"), pl.col("b"), pl.col("c"))
    .filter(
        (pl.col("a").list.unique().list.len() == 1)
        & (pl.col("b").list.unique().list.len() == 1)
    )
    .explode(["a", "b", "c"])
    .select("a", "b", "c")
)

执行结果：

shape: (2, 3)
┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ str │
╞═════╪═════╪═════╡
│ 4   ┆ 2   ┆ b   │
│ 4   ┆ 2   ┆ f   │
└─────┴─────┴─────┘

更直接优雅的实现方案 #

利用Polars的窗口函数，直接在原数据行上计算过滤条件，避免分组聚合再爆炸的冗余步骤：

方案一：用is_unique().over()直接判断 #

df = pl.DataFrame(data).with_columns(
    sum_ab=pl.col("a") + pl.col("b")
).filter(
    pl.col("a").is_unique().over("sum_ab") 
    & pl.col("b").is_unique().over("sum_ab")
).drop("sum_ab")

方案二：统计唯一值数量后过滤 #

如果需要更明确的统计逻辑，可先计算每个sum_ab分组下a和b的唯一值数量，再过滤：

df = pl.DataFrame(data).with_columns(
    sum_ab=pl.col("a") + pl.col("b"),
    # 统计每个sum_ab分组下a的唯一值数量
    a_unique_cnt=pl.col("a").unique().len().over("sum_ab"),
    # 统计每个sum_ab分组下b的唯一值数量
    b_unique_cnt=pl.col("b").unique().len().over("sum_ab")
).filter(
    pl.col("a_unique_cnt") == 1 
    & pl.col("b_unique_cnt") == 1
).drop("sum_ab", "a_unique_cnt", "b_unique_cnt")

优化优势 #

• 逻辑更直观：直接在原数据行上关联分组统计结果，无需对数据进行聚合-爆炸的转换
• 代码更简洁：减少中间步骤嵌套，可读性更强
• 性能更高效：避免数据结构转换带来的额外开销

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