基于列中列表项匹配的高效DataFrame行合并算法需求

你的问题本质上是在解决图的连通分量问题：每一行相当于一个节点，只要两行的b列表存在交集，这两个节点就属于同一个连通分量。原O(n²)的实现对百万级数据来说完全无法承受，而用**并查集（Union-Find/DSU）**可以把时间复杂度降到近乎线性，完美匹配你的性能需求。

核心优化思路 #

1. 构建元素-行索引映射：记录每个b中的元素对应的所有行索引，快速定位有共同元素的关联行。
2. 并查集合并连通分量：利用并查集的高效find/union操作（路径压缩+按秩合并优化后，操作复杂度接近O(1)），把所有连通的行归为同一组。
3. 按组合并结果：将同一连通组内的a值去重合并为列表，b的所有元素去重后合并为列表。

Python实现代码 #

import pandas as pd
from collections import defaultdict

class UnionFind:
    def __init__(self, size):
        self.parent = list(range(size))
        self.rank = [0]*size
    
    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:
            return
        if self.rank[x_root] < self.rank[y_root]:
            self.parent[x_root] = y_root
        else:
            self.parent[y_root] = x_root
            if self.rank[x_root] == self.rank[y_root]:
                self.rank[x_root] += 1

def merge_rows_optimized(df):
    n = len(df)
    if n == 0:
        return pd.DataFrame(columns=df.columns)
    
    # 步骤1：建立元素到行索引的映射
    elem_to_indices = defaultdict(list)
    for idx, b_list in enumerate(df['b']):
        for elem in b_list:
            elem_to_indices[elem].append(idx)
    
    # 步骤2：用并查集合并所有连通的行
    uf = UnionFind(n)
    for indices in elem_to_indices.values():
        if len(indices) <= 1:
            continue
        # 将当前元素关联的所有行合并到同一个连通分量
        root_idx = indices[0]
        for idx in indices[1:]:
            uf.union(root_idx, idx)
    
    # 步骤3：按连通分量分组，合并a和b的值
    groups = defaultdict(lambda: {'a': set(), 'b': set()})
    for idx in range(n):
        root = uf.find(idx)
        groups[root]['a'].add(df['a'].iloc[idx])
        groups[root]['b'].update(df['b'].iloc[idx])
    
    # 转换为目标DataFrame格式
    merged_data = []
    for group in groups.values():
        merged_data.append({
            'a': list(group['a']),
            'b': list(group['b'])
        })
    
    return pd.DataFrame(merged_data, columns=df.columns)

# 测试示例1
df1 = pd.DataFrame({'a': ['1', '2', '3', '4', '5'], 'b': [['a', 'b', 'e'], ['a', 'g'], ['c', 'f'], ['d'], ['b']]})
df1_merged = merge_rows_optimized(df1)
print('Original DF 1:')
print(df1.to_string())
print('Merged DF 1:')
print(df1_merged.to_string())

# 测试示例2
df2 = pd.DataFrame({'a': ['1', '3', '4', '6', '9'], 'b': [['a', 'b', 'e'], ['a', 'g', 'f'], ['c', 'f'], ['d', 'h'], ['b', 'g', 'h']]})
df2_merged = merge_rows_optimized(df2)
print('\nOriginal DF 2:')
print(df2.to_string())
print('Merged DF 2:')
print(df2_merged.to_string())

性能分析 #

• 时间复杂度：O(M α(N))，其中M是所有b列表的总元素数，N是DataFrame行数。α是阿克曼函数的反函数，增长极慢，实际场景中可视为常数，这个复杂度完全能支撑百万级数据的处理。
• 空间复杂度：O(M + N)，主要用于存储元素-索引映射和并查集结构。

扩展优化方向 #

如果数据量极端庞大（比如b总元素数超千万），可以考虑：

1. 并行构建元素映射：用multiprocessing或Dask分块处理DataFrame的行，再合并映射结果。
2. 跨语言性能提升：用C++实现核心的并查集和元素映射逻辑，通过pybind11封装给Python调用；或用Java的HashMap+并查集实现，处理速度会比纯Python更快。

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