数据框中列表列的极坐标交集

Question

import polars as pl

df = pl.DataFrame({'a': [[1, 2, 3], [8, 9, 4]], 'b': [[2, 3, 4], [4, 5, 6]]})

因此给定数据框 df

    a           b
[1, 2, 3]   [2, 3, 4]
[8, 9, 4]   [4, 5, 6]

我想得到c列，它是a和b的交集

    a           b          c
[1, 2, 3]   [2, 3, 4]    [2, 3]
[8, 9, 4]   [4, 5, 6]     [4]

我知道我可以将 apply 函数与 python 集合交集一起使用，但我想使用极坐标表达式来实现。

Answer 1

极坐标 >= 0.18.10

对列表使用集合操作：

df.select(
   intersection = pl.col('a').list.set_intersection('b'),
   difference = pl.col('a').list.set_difference('b'),
   union = pl.col('a').list.set_union('b')
)  

极坐标 >= 0.18.5，极坐标 < 0.18.10

对列表使用集合操作（使用旧名称）：

df.select(
   intersection = pl.col('a').list.intersection('b'),
   difference = pl.col('a').list.difference('b'),
   union = pl.col('a').list.union('b')
)  

极地 < 0.18.5

我们可以使用

arr.eval

表达式来完成交集。

arr.eval

表达式允许我们将列表视为系列/列，以便我们可以使用与列和系列相同的上下文和表达式。

首先，让我们扩展您的示例，以便我们可以展示当交集为空时会发生什么。

df = pl.DataFrame(
    {
        "a": [[1, 2, 3], [8, 9, 4], [0, 1, 2]],
        "b": [[2, 3, 4], [4, 5, 6], [10, 11, 12]],
    }
)
df

shape: (3, 2)
┌───────────┬──────────────┐
│ a         ┆ b            │
│ ---       ┆ ---          │
│ list[i64] ┆ list[i64]    │
╞═══════════╪══════════════╡
│ [1, 2, 3] ┆ [2, 3, 4]    │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [8, 9, 4] ┆ [4, 5, 6]    │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [0, 1, 2] ┆ [10, 11, 12] │
└───────────┴──────────────┘

算法

有两种方法可以实现这一点。 第一个可扩展到两个以上集合的交集（请参阅下面的其他注释）。

df.with_column(
    pl.col("a")
    .arr.concat('b')
    .arr.eval(pl.element().filter(pl.count().over(pl.element()) == 2))
    .arr.unique()
    .alias('intersection')
)

或

df.with_column(
    pl.col("a")
    .arr.concat('b')
    .arr.eval(pl.element().filter(pl.element().is_duplicated()))
    .arr.unique()
    .alias('intersection')
)

shape: (3, 3)
┌───────────┬──────────────┬──────────────┐
│ a         ┆ b            ┆ intersection │
│ ---       ┆ ---          ┆ ---          │
│ list[i64] ┆ list[i64]    ┆ list[i64]    │
╞═══════════╪══════════════╪══════════════╡
│ [1, 2, 3] ┆ [2, 3, 4]    ┆ [2, 3]       │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [8, 9, 4] ┆ [4, 5, 6]    ┆ [4]          │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [0, 1, 2] ┆ [10, 11, 12] ┆ []           │
└───────────┴──────────────┴──────────────┘

如何运作

我们首先将两个列表连接成一个列表。 两个列表中的任何元素都会出现两次。

df.with_column(
    pl.col("a")
    .arr.concat('b')
    .alias('ablist')
)

shape: (3, 3)
┌───────────┬──────────────┬────────────────┐
│ a         ┆ b            ┆ ablist         │
│ ---       ┆ ---          ┆ ---            │
│ list[i64] ┆ list[i64]    ┆ list[i64]      │
╞═══════════╪══════════════╪════════════════╡
│ [1, 2, 3] ┆ [2, 3, 4]    ┆ [1, 2, ... 4]  │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [8, 9, 4] ┆ [4, 5, 6]    ┆ [8, 9, ... 6]  │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [0, 1, 2] ┆ [10, 11, 12] ┆ [0, 1, ... 12] │
└───────────┴──────────────┴────────────────┘

然后我们可以使用

arr.eval

函数，它允许我们将连接列表视为一个系列/列。 在本例中，我们将使用

filter

上下文来查找出现多次的任何元素。 （列表上下文中的

polars.element

表达式的使用方式类似于系列中使用的

polars.col

。）

df.with_column(
    pl.col("a")
    .arr.concat('b')
    .arr.eval(pl.element().filter(pl.count().over(pl.element()) == 2))
    .alias('filtered')
)

shape: (3, 3)
┌───────────┬──────────────┬───────────────┐
│ a         ┆ b            ┆ filtered      │
│ ---       ┆ ---          ┆ ---           │
│ list[i64] ┆ list[i64]    ┆ list[i64]     │
╞═══════════╪══════════════╪═══════════════╡
│ [1, 2, 3] ┆ [2, 3, 4]    ┆ [2, 3, ... 3] │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [8, 9, 4] ┆ [4, 5, 6]    ┆ [4, 4]        │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [0, 1, 2] ┆ [10, 11, 12] ┆ []            │
└───────────┴──────────────┴───────────────┘

注意：上述步骤也可以使用

is_duplicated

表达式来表示。 （在“其他注释”部分中，我们将看到在计算两个以上集合的交集时，使用

is_duplicated

将不起作用。）

df.with_column(
    pl.col("a")
    .arr.concat('b')
    .arr.eval(pl.element().filter(pl.element().is_duplicated()))
    .alias('filtered')
)

剩下的就是使用

arr.unique

表达式（即开头所示的结果）从结果中删除重复项。

其他注意事项

我假设您的列表确实是集合，因为元素在每个列表中仅出现一次。 如果原始列表中有重复项，我们可以在串联步骤之前将

arr.unique

应用于每个列表。

此外，这个过程可以扩展以查找两个以上集合的交集。 只需将所有列表连接在一起，然后将

filter

步骤从

== 2

更改为

== n

（其中

n

是组数）。 （注意：使用上面的

is_duplicated

表达式不适用于两个以上的集合。）

arr.eval

方法确实有一个

parallel

关键字。 您可以尝试将其设置为

True

，看看它在您的特定情况下是否会产生更好的性能。

其他设置操作

对称差异：将

filter

标准更改为

== 1

（并省略

arr.unique

步骤。）

并集：使用

arr.concat

，然后使用

arr.unique

。

设置差异：计算交集（如上所述），然后连接原始列表/集合并过滤仅出现一次的项目。 或者，对于较小的列表大小，您可以将“a”连接到自身，然后连接到“b”，然后过滤出现两次（但不是三次）的元素。