Python代码优化求助：百万级输入下嵌套循环超时问题

问题说明 #

输入规则如下：

• 首行输入正整数X（X≥0），后续X行每行两个单词，表示这两个单词必须同组；
• 接下来输入正整数Y（Y≥0），后续Y行每行两个单词，表示这两个单词不能同组；
• 最后输入正整数G（G≥1），后续G行每行三个不同单词，表示这三个单词被分到同一组。

需要输出被违反的约束总数，结果范围在0到X+Y之间。

原代码 #

mustConstraint = set()
notConstraint = set()
violated = 0
satisfied = 0 

for i in range(0, int(input(''))):
    constraint = input('')
    mustConstraint.add(frozenset(constraint.split()))

for i in range(0, int(input(''))):
    constraint = input('')
    notConstraint.add(frozenset(constraint.split()))

for i in range(0, int(input(''))):
    group = input('')
    group = set(group.split())

    for x in mustConstraint:
        if x & group == x:
            satisfied +=1
        
    for y in notConstraint:
        if y & group == y:
            violated += 1

violated += len(mustConstraint) - satisfied

print(violated)

性能瓶颈 #

当G达到30万量级时，原代码的嵌套循环会带来**O(G*(X+Y))**的时间复杂度，极端场景下（如X/Y各10万）会产生200亿次迭代，完全无法在4秒时限内完成。核心问题是对每个组遍历所有约束做检查，效率极低。

优化方案 #

核心思路：反向统计，从约束入手而非组 #

不需要对每个组遍历所有约束，而是先建立「单词→所属组」的映射，再直接检查每个约束是否被满足/违反：

1. 用字典存储每个单词对应的组标识（比如组的索引）；
2. 遍历必须约束：若两个单词不在同一组，计数一次违反；
3. 遍历禁止约束：若两个单词在同一组，计数一次违反。

该方案时间复杂度降为O(X+Y+G)，完全适配大规模输入。

优化后代码 #

# 读取必须约束
x = int(input())
must_constraints = []
for _ in range(x):
    a, b = input().split()
    must_constraints.append((a, b))

# 读取禁止约束
y = int(input())
not_constraints = []
for _ in range(y):
    a, b = input().split()
    not_constraints.append((a, b))

# 建立单词到组的映射：用组的索引作为唯一标识
word_to_group = {}
g = int(input())
for group_idx in range(g):
    words = input().split()
    for word in words:
        word_to_group[word] = group_idx

violated = 0

# 检查必须约束：不在同一组则违反
for a, b in must_constraints:
    if word_to_group[a] != word_to_group[b]:
        violated += 1

# 检查禁止约束：在同一组则违反
for a, b in not_constraints:
    if word_to_group[a] == word_to_group[b]:
        violated += 1

print(violated)

优化细节 #

• 放弃集合存储约束，直接保留原始单词对，减少集合操作开销；
• 字典查询单词所属组的时间复杂度为O(1)，大幅提升查询效率；
• 完全消除嵌套循环，所有操作均为线性遍历，性能显著提升；
• 若题目保证约束中的单词必然出现在组中，可直接使用word_to_group[a]，无需get方法（避免KeyError）。

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