如何在Python中按链式公共元素对嵌套列表进行分组？

问题描述 #

给定一个嵌套列表，需按以下规则分组：

• 组内至少有两个子列表共享至少一个公共元素
• 组内每个子列表无需与所有子列表有公共元素，只需能与组内另一个子列表形成链式关联即可

示例：

输入: f(ls = [[3, 4, 2], [1, 3, 5], [9, 8, 7], [8, 10, 34], [34], [1], [23, 22, 21], [22], [21, 29], [100]])

输出:
{0: [[3, 4, 2], [1, 3, 5], [1]], 1: [[9, 8, 7], [8, 10, 34], [34]], 2: [[23, 22, 21], [22], [21, 29]], 3: [[100]]}

解决思路 #

这个问题本质是找连通分量：把每个子列表看作一个节点，只要两个子列表有公共元素，就认为它们连通；通过链式关联连通的节点会被归为同一组。可以用**并查集（Union-Find）**数据结构高效处理这类连通性问题，步骤如下：

1. 初始化并查集，每个子列表单独作为一个集合
2. 建立元素到子列表索引的映射：记录每个元素出现在哪些子列表里
3. 遍历每个元素对应的索引列表，将这些索引全部合并到同一个集合
4. 最后把属于同一集合的子列表分组，整理成要求的字典格式

代码实现 #

class UnionFind:
    def __init__(self, size):
        self.parent = list(range(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):
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x != root_y:
            self.parent[root_y] = root_x

def f(ls):
    n = len(ls)
    if n == 0:
        return {}
    
    # 初始化并查集
    uf = UnionFind(n)
    
    # 构建元素到子列表索引的映射
    elem_to_indices = {}
    for idx, sublist in enumerate(ls):
        for num in sublist:
            if num not in elem_to_indices:
                elem_to_indices[num] = []
            elem_to_indices[num].append(idx)
    
    # 合并所有共享同一元素的子列表索引
    for indices in elem_to_indices.values():
        if len(indices) >= 2:
            first_idx = indices[0]
            for idx in indices[1:]:
                uf.union(first_idx, idx)
    
    # 按连通分量分组
    groups = {}
    for idx in range(n):
        root = uf.find(idx)
        if root not in groups:
            groups[root] = []
        groups[root].append(ls[idx])
    
    # 将组的键转换为连续的0开始的整数，匹配示例输出格式
    return {i: group for i, group in enumerate(groups.values())}

# 测试示例
ls = [[3, 4, 2], [1, 3, 5], [9, 8, 7], [8, 10, 34], [34], [1], [23, 22, 21], [22], [21, 29], [100]]
print(f(ls))

代码说明 #

• UnionFind类：实现了路径压缩的并查集，保证查找和合并操作的时间复杂度接近O(1)
• 元素-索引映射：快速定位所有包含同一元素的子列表，避免两两比较子列表的时间浪费
• 分组整理：最后将同一根节点的子列表归为一组，并且把组的键转换为连续整数，和示例输出格式完全匹配

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