Python中如何合并字典内有交集的子列表并关联记录ID？

问题描述 #

现有如下字典：

data = {
    'similar_record_id': ['A-1','A-2','A-3','A-4','A-5'],
    'idx_pair' : [[25, 26],[835, 836],[834, 836, 835],[67, 69, 68],[62, 68, 66, 69, 67, 65]]
}

需合并idx_pair中的子列表：任意两个存在共同元素的子列表必须合并（例如[835,836]与[834,836,835]合并为[834,835,836]），同时关联对应的similar_record_id，最终得到如下结构的结果：

{
    'similar_record_id': ['A-1','A-2','A-4'],
    'idx_pair' : [[25, 26],[834,835, 836],[62, 65, 66, 67, 68, 69]]
}

注：原预期输出中similar_record_id的第三个元素应为A-4（对应原数据中A-4、A-5的子列表合并组），属于笔误修正。

解决方案：使用并查集（Union-Find） #

并查集是处理连通分量合并问题的高效工具，以下是实现代码：

并查集类实现 #

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

数据处理逻辑 #

data = {
    'similar_record_id': ['A-1','A-2','A-3','A-4','A-5'],
    'idx_pair' : [[25, 26],[835, 836],[834, 836, 835],[67, 69, 68],[62, 68, 66, 69, 67, 65]]
}

# 初始化并查集实例
uf = UnionFind()

# 为所有元素初始化父节点
for idx_list in data['idx_pair']:
    for num in idx_list:
        if num not in uf.parent:
            uf.parent[num] = num

# 合并每个子列表内的元素，确保同一子列表的元素属于同一集合
for idx_list in data['idx_pair']:
    if len(idx_list) >= 2:
        first_num = idx_list[0]
        for num in idx_list[1:]:
            uf.union(first_num, num)

# 按根节点分组元素，得到合并后的元素列表
root_to_elements = {}
for num in uf.parent:
    root = uf.find(num)
    if root not in root_to_elements:
        root_to_elements[root] = []
    root_to_elements[root].append(num)

# 记录每个连通分量首次出现时对应的similar_record_id，保证结果顺序与原数据一致
root_to_first_id = {}
seen_roots = set()
root_order = []

for idx, idx_list in enumerate(data['idx_pair']):
    if not idx_list:
        continue
    current_root = uf.find(idx_list[0])
    if current_root not in seen_roots:
        seen_roots.add(current_root)
        root_order.append(current_root)
        root_to_first_id[current_root] = data['similar_record_id'][idx]

# 构造最终结果
final_result = {
    'similar_record_id': [root_to_first_id[root] for root in root_order],
    'idx_pair': [sorted(list(set(root_to_elements[root]))) for root in root_order]
}

# 打印结果
print(final_result)

代码说明 #

1. 并查集核心：通过find方法实现路径压缩，union方法合并集合，确保高效处理大规模数据。
2. 元素初始化：遍历所有子列表元素，为每个元素建立独立集合。
3. 子列表合并：将每个子列表内的元素合并到同一集合，标记它们属于同一连通分量。
4. 分组与排序：按根节点将元素分组，去重并排序得到整洁的合并后列表。
5. ID关联：按原数据中连通分量首次出现的顺序，关联对应的similar_record_id，保证结果顺序符合直觉。

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