大规模唯一元素集合交集关联映射构建算法求解问询

问题回顾 #

给定：

• 元素集合S（如非负整数集合）
• 映射N：将S中部分元素映射到S的子集，每个子集元素唯一有序

示例输入：

S = {0, 1, 2, 3, 4, 5, 6, 7, 8}
N = {
   0: [1, 2, 3],
   4: [2, 3, 5, 6],
   6: [5, 7, 8],
}

目标：构建映射M，其中每个N的键s_j对应N的键子集S_j，当且仅当N[s_j]与N[s]存在公共元素时，s被纳入S_j。示例结果：

{
   0: [4],
   4: [0, 6],
   6: [4],
}

参数规模：

• N的键数量约10万
• 每个N[s]的元素数量约10
• S的唯一元素总数约10万

代码正确性分析 #

你的倒排索引思路是正确的：

1. 先记录每个元素x被哪些N的键包含（即哪些s满足x∈N[s]）
2. 对每个元素对应的键集合，将集合内所有两两不同的键互相加入对方的相交集合

最终能得到符合要求的映射M，逻辑上是正确的。但代码存在严重的效率瓶颈，无法处理10万规模的输入。

时间复杂度分析 #

你认为时间复杂度是O(|N|)，这是错误的，当前代码的复杂度远高于此：

• 瓶颈步骤：构建inclusion_by_index时，对每个全局元素x（最多10万），遍历所有N的键（10万）并检查x是否在其邻居集合中。这部分的时间复杂度是O(KM)，其中K是N的键数，M是全局元素数，计算得1e51e5=1e10次操作——这完全超出了常规程序的运行时间上限。
• 其余步骤：product遍历部分的复杂度是Σ(m_x²)，其中m_x是元素x对应的键数量。在平均情况下（每个元素被约10个键包含），Σ(m_x²)=1e5*(10²)=1e7次操作，这部分是可接受的。

优化方案 #

核心优化：高效构建倒排索引 #

将构建inclusion_by_index的方式从"遍历元素→遍历所有键"改为"遍历键→遍历其元素"，这样时间复杂度降为O(total_elements)（总元素数约1e6）：

from collections import defaultdict
from itertools import product
from functools import reduce
from operator import or_
from typing import Sequence, Set, Union, List, Mapping
import numpy as np

def merge_sets(sets: Sequence[Set]) -> Set:
    """合并多个集合为一个"""
    if not sets:
        return set()
    return reduce(or_, sets)

def neighbors_to_mutexes(
        indexes: Union[List[int], np.ndarray],
        neighbors_by_index: Union[List[List[int]], np.ndarray],
) -> Mapping[int, Set[int]]:
    """找出邻居列表存在交集的索引对"""
    inclusion_by_index = defaultdict(set)  # 元素x对应的所有包含它的键
    
    # 优化：遍历每个键，将其加入对应元素的集合
    for idx, neighbors in zip(indexes, neighbors_by_index):
        for elem in neighbors:
            inclusion_by_index[elem].add(idx)
    
    intersecting = defaultdict(set)
    # 可选优化：只处理i<j的配对，避免重复操作
    for including_set in inclusion_by_index.values():
        lst = list(including_set)
        for i in range(len(lst)):
            for j in range(i+1, len(lst)):
                a, b = lst[i], lst[j]
                intersecting[a].add(b)
                intersecting[b].add(a)
    
    return intersecting

其他小优化 #

• 省略将neighbors_by_index转为集合的操作：优化后不需要检查元素存在性，无需额外转换
• 配对时只处理i<j的组合：减少一半操作次数，避免重复添加相同配对

优化后的时间复杂度 #

• 倒排索引构建：O(total_elements) = 1e6次操作
• 配对处理：Σ(m_x*(m_x-1)/2)*2 = Σ(m_x² - m_x) ≈1e7次操作
• 整体复杂度为O(total_elements + Σ(m_x²))，完全符合10万规模的性能要求

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