如何按字符串子串对Polars DataFrame分组聚类？具体场景答疑

云上老周

2026-6-26

问题解答

原始DataFrame

import polars as pl

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

1. 如何让id为1、2、3的行处于同一组？

观察values列的关联逻辑：id1对应"A"，id2对应"B"，id3的"A--B"直接关联了A和B，因此需要基于元素的连通性做聚类。可以结合Polars的数据处理能力和图论工具实现：

具体步骤：

拆分values列的复合元素，将"A--B"这类条目拆分为独立元素并展开成行
构建图结构：给同一id下的元素两两连边，让关联元素处于同一连通分量
计算每个元素的连通分量标签，再映射回原DataFrame，同一连通分量的id自动归为一组

实现代码：

import networkx as nx

# 拆分values并展开为多行
split_df = df.with_columns(pl.col("values").str.split("--").alias("elements")).explode("elements")

# 构建图，关联同一id内的元素
G = nx.Graph()
G.add_nodes_from(split_df["elements"].unique())

for _, group in split_df.group_by("id"):
    elements = group["elements"].to_list()
    # 同一id内的元素两两连边，确保连通
    if len(elements) >= 2:
        for i in range(len(elements)):
            for j in range(i+1, len(elements)):
                G.add_edge(elements[i], elements[j])

# 生成元素到分组id的映射
component_map = {}
for group_idx, component in enumerate(nx.connected_components(G)):
    for elem in component:
        component_map[elem] = group_idx

# 映射回原DataFrame，得到每个id的分组
result = split_df.with_columns(pl.col("elements").map_dict(component_map).alias("group_id"))\
                .group_by("id").agg(pl.col("group_id").first())\
                .join(df, on="id")

print(result)

执行后，id1、2、3的group_id会是同一个值，满足同组需求。