USACO Silver Wormhole Sort问题Python解法超时优化求助

你的代码逻辑正确，但在大数据测试用例上超时，核心原因是递归DFS的开销过大，以及每次检查连通性时的遍历效率不足。以下是针对性的优化方案：

核心优化思路 #

用**并查集（Disjoint Set Union, DSU）**替代DFS处理连通性问题。并查集的合并与查询操作时间复杂度接近O(α(n))（α为阿克曼函数的反函数，增长极慢），远优于DFS的O(n+m)每次检查，尤其适合二分查找中多次调用连通性检查的场景。

优化后的代码 #

import sys

def main():
    fin = open("wormsort.in", "r")
    fout = open("wormsort.out", "w")
    
    n, m = map(int, fin.readline().split())
    cows = list(map(int, fin.readline().split()))
    
    # 提前判断是否已排序
    sorted_cows = list(range(1, n+1))
    if cows == sorted_cows:
        fout.write("-1\n")
        return
    
    # 存储所有虫洞边
    edges = []
    for _ in range(m):
        a, b, w = map(int, fin.readline().split())
        edges.append((w, a, b))
    
    # 按权重排序，用于二分查找
    edges.sort()
    weights = [w for w, _, _ in edges]
    
    class DSU:
        def __init__(self, size):
            self.parent = list(range(size + 1))  # 1-based索引
            self.rank = [0] * (size + 1)
        
        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:
                return
            # 按秩合并
            if self.rank[x_root] < self.rank[y_root]:
                self.parent[x_root] = y_root
            else:
                self.parent[y_root] = x_root
                if self.rank[x_root] == self.rank[y_root]:
                    self.rank[x_root] += 1
    
    def is_possible(min_weight):
        dsu = DSU(n)
        # 合并所有权重 >= min_weight 的边
        for w, a, b in edges:
            if w >= min_weight:
                dsu.union(a, b)
            else:
                # 因为边已排序，后面的权重更小，直接跳出
                break
        # 检查每个位置i和对应奶牛应在的位置是否连通
        for i in range(n):
            pos = i + 1
            cow = cows[i]
            if dsu.find(pos) != dsu.find(cow):
                return False
        return True
    
    # 二分查找最大满足条件的权重
    low = 0
    high = m - 1
    ans = 0
    while low <= high:
        mid = (low + high) // 2
        if is_possible(weights[mid]):
            ans = weights[mid]
            low = mid + 1
        else:
            high = mid - 1
    
    fout.write(f"{ans}\n")

if __name__ == "__main__":
    main()

关键优化点说明 #

1. 替换DFS为并查集：

   • 去掉递归调用，避免Python递归的函数调用开销和栈溢出风险（即使设置了递归深度，递归效率仍远低于并查集）。
   • 并查集的路径压缩和按秩合并保证了近乎常数时间的操作效率。

2. 边的预处理优化：

   • 直接存储所有边并按权重排序，无需维护邻接表，减少内存占用和遍历开销。
   • 在is_possible函数中，利用边已排序的特性，权重小于当前阈值时直接跳出循环，减少不必要的合并操作。

3. 二分逻辑调整：

   • 优化二分的边界处理，确保找到最大的满足条件的权重，避免原代码中边界计算的误差。

4. 索引清晰化：

   • 统一使用1-based索引处理奶牛位置，避免原代码中索引转换的混乱，减少出错概率。

额外性能提升建议 #

如果想进一步优化，可以尝试贪心算法：

• 将边按权重降序排序，依次合并边，同时维护需要连通的位置对（i和cows[i-1]）。每次合并后，检查所有未连通的对是否已连通，当所有对都连通时，当前边的权重就是答案。这种方法避免了二分查找的多次检查，时间复杂度为O(mlogm + nα(n))，在某些场景下效率更高。

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