如何高效合并存在元素交集的含数组类型Python DataFrame?
高效合并两个含数组的DataFrame(百万级数据量适配)
针对你需要合并两个DataFrame中存在至少一个共同元素的数组的需求,结合百万级数据量的性能要求,这里推荐基于**哈希表+并查集(Union-Find)**的实现方案,时间复杂度接近线性,能高效处理大规模数据。
先明确输入数据结构
先把你提供的示例DataFrame用更清晰的方式展示:
DataFrame 1
| 索引 | 数组 |
|---|---|
| 0 | [1, 50] |
| 1 | [60, 61, 62, 63, 66, 64, 65, 67] |
| 2 | [7, 8, 9] |
| 3 | [80, 81, 72, 83] |
| 4 | [90, 91, 92] |
| 5 | [20, 21, 22, 23, 24, 25, 26, 27, 28] |
| 6 | [200, 201] |
| 7 | [301, 300] |
DataFrame 2
| 索引 | 数组 |
|---|---|
| 0 | [1, 2] |
| 1 | [3, 4] |
| 2 | [5, 6] |
| 3 | [7, 71, 72, 73, 74, 75, 76] |
| 4 | [10, 11, 12] |
| 5 | [100, 100, 102] |
| 6 | [30, 31] |
| 7 | [40, 41] |
核心实现思路(高效处理百万级数据)
因为数据量达百万级,绝对不能用嵌套循环遍历所有数组对(时间复杂度O(M*N),M和N是两个DataFrame的数组数量,完全不可行)。我们需要用以下两步实现:
1. 构建元素到数组代表的映射
首先,给每个数组分配一个唯一标识(比如可以用原DataFrame的索引加上来源标识,如df1_0、df2_3),然后遍历所有数组中的元素,记录每个元素对应的第一个遇到的数组标识。
2. 用并查集合并连通的数组
遍历每个数组,对于数组中的每个元素,找到该元素对应的数组标识,将当前数组标识与这个标识合并(在并查集中)。这样,所有有共同元素的数组会被归到同一个连通分量中。
3. 合并每个连通分量的数组
最后,遍历所有数组,将同一连通分量中的数组合并(去重,可选),生成最终结果。
代码实现(Python)
下面是具体的Python代码实现,用到pandas处理DataFrame,以及自定义的并查集:
import pandas as pd from collections import defaultdict # 自定义并查集类,带路径压缩优化 class UnionFind: def __init__(self): self.parent = {} 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 # 示例DataFrame构建 df1 = pd.DataFrame({ 'array': [ [1, 50], [60, 61, 62, 63, 66, 64, 65, 67], [7, 8, 9], [80, 81, 72, 83], [90, 91, 92], [20, 21, 22, 23, 24, 25, 26, 27, 28], [200, 201], [301, 300] ]}, index=[0,1,2,3,4,5,6,7]) df2 = pd.DataFrame({ 'array': [ [1, 2], [3, 4], [5, 6], [7, 71, 72, 73, 74, 75, 76], [10, 11, 12], [100, 100, 102], [30, 31], [40, 41] ]}, index=[0,1,2,3,4,5,6,7]) # 步骤1:收集所有数组并分配唯一标识 all_arrays = [] # 标识格式:来源_df_索引,确保每个数组的标识唯一 for idx, arr in df1['array'].items(): all_arrays.append(('df1_' + str(idx), arr)) for idx, arr in df2['array'].items(): all_arrays.append(('df2_' + str(idx), arr)) # 步骤2:建立元素到数组标识的映射,初始化并查集 element_to_group = {} uf = UnionFind() for group_id, arr in all_arrays: # 初始化当前数组的并查集父节点 if group_id not in uf.parent: uf.parent[group_id] = group_id # 遍历数组中的每个元素,处理关联关系 for num in arr: if num in element_to_group: # 元素已关联其他数组,合并两个数组的组 uf.union(group_id, element_to_group[num]) else: # 元素首次出现,关联到当前数组的组 element_to_group[num] = group_id # 步骤3:合并同一连通分量的数组 group_to_arrays = defaultdict(list) for group_id, arr in all_arrays: root = uf.find(group_id) group_to_arrays[root].extend(arr) # 可选:去重并排序(根据业务需求决定是否需要) result = {root: sorted(list(set(arr))) for root, arr in group_to_arrays.items()} # 转换回DataFrame格式(也可以直接用字典/列表作为结果) result_df = pd.DataFrame({'merged_array': list(result.values())}, index=list(result.keys())) print(result_df)
性能说明
- 时间复杂度:O(T α(T)),其中T是所有数组中元素的总数量,α是阿克曼函数的反函数(几乎可以看作常数),完全适配百万级数据量。
- 空间复杂度:O(T),主要用于存储元素到组的映射和并查集的父节点。
关键优化点
- 避免了嵌套循环,用哈希表快速查找元素对应的组,彻底规避了O(M*N)的低效遍历。
- 并查集的路径压缩优化保证了合并和查找操作的时间复杂度接近常数,即使数据量极大也能快速处理。
内容的提问来源于stack exchange,提问作者Supun
相关产品推荐
相关产品推荐

