USACO 2020 US Open铜组第3题Python代码超时优化求助

核心瓶颈分析 #

原代码时间复杂度为O(n*t²)，当n=100、t=250时理论运算量达625万次，叠加逻辑错误、无效计算、低效数据结构的开销，最终导致测试用例13超时。

可落地的优化方案 #

• 前置剪枝无效初始病牛：只有最终状态为患病（end_cows[p0] == 1）的牛才可能是初始病牛，直接过滤掉90%以上的无效循环，测试用例13仅4个候选初始病牛，循环量直接降为原有的1/25。
• 修正max_k计算逻辑：原代码无论模拟是否成功都更新max_k属于逻辑错误，仅当模拟结果与最终状态一致时才需要更新min_k和max_k。
• 减少k的遍历次数：对每个候选初始病牛p0，只需找到最小的有效k值，再直接测试k=t的情况：如果k=t时也有效，说明所有大于等于该最小k的取值都有效，无需遍历中间所有k值，可直接标记存在无穷大的k。
• 优化数据结构与操作：
  • 用集合存储possible候选初始病牛，成员判断时间复杂度从O(n)降至O(1)
  • 用元组代替自定义Shake类存储握手记录，属性访问速度更快
  • 避免使用全局变量，函数参数/闭包访问的效率更高
• 提前终止无效模拟：模拟过程中如果出现和最终状态矛盾的情况（比如某头牛最终患病但模拟过程中已经不可能被感染），可提前终止当前模拟，节省时间。

优化后参考代码 #

def main():
    import sys
    input = sys.stdin.read().splitlines()
    ptr = 0
    n, t = map(int, input[ptr].split())
    ptr +=1
    end_cows = [int(c) for c in input[ptr].strip()]
    ptr +=1
    shakes = []
    for _ in range(t):
        time, a, b = map(int, input[ptr].split())
        shakes.append((time, a-1, b-1))
        ptr +=1
    # 按时间排序握手
    shakes.sort()
    
    # 预筛选可能的初始病牛
    candidates = [i for i in range(n) if end_cows[i] == 1]
    possible = set()
    min_k = t
    has_infinite = False
    
    def sim(p0, k):
        infected = [False]*n
        infected[p0] = True
        remain = [0]*n
        remain[p0] = k
        for _, a, b in shakes:
            # 先记录双方是否能传播
            a_can_spread = infected[a] and remain[a] > 0
            b_can_spread = infected[b] and remain[b] > 0
            # 传播
            if a_can_spread:
                infected[b] = True
            if b_can_spread:
                infected[a] = True
            # 扣减次数
            if infected[a]:
                remain[a] -= 1
            if infected[b]:
                remain[b] -= 1
        # 转换为0/1列表对比
        res = [1 if x else 0 for x in infected]
        return res == end_cows
    
    for p0 in candidates:
        # 找最小的有效k
        p0_min_k = None
        for k in range(t+1):
            if sim(p0, k):
                p0_min_k = k
                break
        if p0_min_k is None:
            continue
        # 加入可能的初始病牛
        possible.add(p0)
        # 更新全局最小k
        if p0_min_k < min_k:
            min_k = p0_min_k
        # 测试k=t是否有效，有效则说明存在无穷大k
        if sim(p0, t):
            has_infinite = True
    
    max_k = "Infinity" if has_infinite else 0
    if not has_infinite:
        # 找最大的有效k
        for p0 in possible:
            for k in range(t, -1, -1):
                if sim(p0, k):
                    if k > max_k:
                        max_k = k
                    break
    
    with open("tracing.out", 'w') as fout:
        fout.write(f"{len(possible)} {min_k} {max_k}\n")

if __name__ == "__main__":
    main()

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