求助：Pandas中基于ID元组元素交集的分组实现方法

问题分析 #

你需要的不是按元组完全匹配分组，而是将存在元素交集的元组归为同一组——这本质是图论中的连通分量问题：把每个唯一元组看作节点，若两个元组有公共元素则连边，最终每个连通分量就是一个分组。

实现方案 #

下面提供两种可行的实现方式，分别基于并查集（高效轻量）和NetworkX图库（简洁直观）。

方式1：并查集算法（高效轻量） #

并查集适合处理大规模数据的连通性问题，无需额外依赖第三方图库。

import pandas as pd

# 1. 创建示例DataFrame
data = {
    'Frame_number': ['Frame 1']*3 + ['Frame 2']*3 + ['Frame 3']*3 + ['Frame 4']*3 + ['Frame 5']*3 + ['Frame 6']*3,
    'Color_ID': [(10,5), (11,8), (12,9)]*3 + [(80,5), (11,8), (12,9)]*2 + [(80,6), (11,8), (12,9)]
}
df = pd.DataFrame(data)

# 2. 实现并查集类
class UnionFind:
    def __init__(self, elements):
        self.parent = {elem: elem for elem in elements}
    
    def find(self, x):
        # 路径压缩，提升查找效率
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def union(self, x, y):
        # 合并两个连通分量
        x_root = self.find(x)
        y_root = self.find(y)
        if x_root != y_root:
            self.parent[y_root] = x_root

# 3. 提取唯一Color_ID并初始化并查集
unique_color_ids = df['Color_ID'].unique()
uf = UnionFind(unique_color_ids)

# 4. 遍历元组对，合并有交集的元组
for i in range(len(unique_color_ids)):
    set_i = set(unique_color_ids[i])
    for j in range(i+1, len(unique_color_ids)):
        set_j = set(unique_color_ids[j])
        if set_i & set_j:  # 存在元素交集
            uf.union(unique_color_ids[i], unique_color_ids[j])

# 5. 生成组ID映射
group_mapping = {}
root_to_group = {}
current_group = 0
for elem in unique_color_ids:
    root = uf.find(elem)
    if root not in root_to_group:
        root_to_group[root] = current_group
        current_group += 1
    group_mapping[elem] = root_to_group[root]

# 6. 将组ID添加到原DataFrame
df['Group_ID'] = df['Color_ID'].map(group_mapping)

# 查看分组结果
print(df.groupby('Group_ID')['Color_ID'].unique())

方式2：NetworkX图库（简洁直观） #

如果不需要考虑极致性能，用NetworkX的连通分量功能可以快速实现逻辑，代码更简洁。

import pandas as pd
import networkx as nx

# 1. 创建示例DataFrame（同方式1）
data = {
    'Frame_number': ['Frame 1']*3 + ['Frame 2']*3 + ['Frame 3']*3 + ['Frame 4']*3 + ['Frame 5']*3 + ['Frame 6']*3,
    'Color_ID': [(10,5), (11,8), (12,9)]*3 + [(80,5), (11,8), (12,9)]*2 + [(80,6), (11,8), (12,9)]
}
df = pd.DataFrame(data)

# 2. 提取唯一Color_ID并构建图
unique_color_ids = df['Color_ID'].unique()
G = nx.Graph()
G.add_nodes_from(unique_color_ids)

# 3. 为有交集的元组添加边
for i in range(len(unique_color_ids)):
    set_i = set(unique_color_ids[i])
    for j in range(i+1, len(unique_color_ids)):
        set_j = set(unique_color_ids[j])
        if set_i & set_j:
            G.add_edge(unique_color_ids[i], unique_color_ids[j])

# 4. 获取连通分量并生成组ID映射
connected_components = list(nx.connected_components(G))
group_mapping = {}
for idx, comp in enumerate(connected_components):
    for elem in comp:
        group_mapping[elem] = idx

# 5. 添加组ID到DataFrame
df['Group_ID'] = df['Color_ID'].map(group_mapping)

# 查看分组结果
print(df.groupby('Group_ID')['Color_ID'].unique())

结果验证 #

两种方式都会得到如下分组：

• Group 0: [(10,5), (80,5), (80,6)]
• Group 1: [(11,8)]
• Group 2: [(12,9)]

完全符合你的需求。

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