Python 组群集按字符串中的子字符串或子字符串中的字符串轮询数据框

发布于04月04日

给出这个Polars DataFrame:

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

1, How can I group/cluster it so that 1,2 and 3 ends up in the same group?
2. Can I even achieve having 4 in the same group/cluster?

推荐答案

假设您希望根据子字符串(以--分隔)合并组，不幸的是，这并不简单.您不能将其矢量化，因为一个组的成员可以链接到另一个组，而另一个组链接到另一个组，等等.

一种 Select 是使用图论来识别连接的组件.你可以用networkx和connected_components来做:

import networkx as nx

G = nx.from_pandas_edgelist(df.with_columns(pl.col('values').str.split('--'))
                              .explode('values'), source='id', target='values')

S = set(df['id'])
mapper = {n: i for i, c in enumerate(nx.connected_components(G)) for n in c&S}
# {1: 0, 2: 0, 3: 0, 4: 0, 5: 1}

out = df.group_by(pl.col('id').replace(mapper).alias('group')).agg(pl.all())

输出:

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

图表: