Bellman-Ford算法负环检测实现出错及TLE问题求助

问题背景 #

用Python实现Bellman-Ford算法检测图中的负环，用于解决CSES的负环检测问题，但遇到错误答案（WA）和超时（TLE）问题。原代码如下：

def bellman_ford(n, edges):
    dist = [float('inf')] * (n + 1)
    dist[1] = 0
    prev = [-1] * (n + 1)
    for _ in range(n - 1):
        for u, v, w in edges:
            if dist[u] != float('inf') and  dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
                prev[v] = u
    for u, v, w in edges:
        if dist[u]!=float('inf') and dist[u] + w < dist[v]:
            cycle = [v]
            prev[v] = u
            while u != v:
                cycle.append(u)
                u = prev[u]
            cycle.append(v)
            return True, cycle[::-1]
    return False, None
n, e = map(int, input().split())
edges = []
for _ in range(e):
    u, v, w = map(int, input().split())
    edges.append((u, v, w))

chk, cycle = bellman_ford(n, edges)
if chk == True:
    print("YES")
    print(*cycle)
else:
    print("NO")

问题分析与修复方案 #

1. 错误答案（WA）：仅检测从节点1可达的负环 #

原代码仅将dist[1]初始化为0，其余节点设为无穷大，这只能检测从节点1可达的负环。但题目要求检测图中任意存在的负环，如果负环不在节点1的可达区域内，代码会漏判，导致WA。

修复：将所有节点的初始距离设为0，相当于添加一个虚拟源点向所有节点连权重为0的边，确保无论负环在图的哪个区域都能被检测到。

2. 错误答案（WA）：环的查找逻辑错误 #

原代码在找到可松弛的边u->v后，直接修改prev[v] = u再回溯找环，这种逻辑可能因prev指针指向问题无法正确追踪环，甚至陷入死循环。

修复：当找到可松弛的边时，从v出发沿prev指针走n步（图最多n个节点，走n步必然进入环），找到环中的一个节点，再从该节点回溯直到回到自身，构建完整环。

3. 超时（TLE）：输入处理效率低 #

Python的input()函数在处理大量输入时速度慢，循环调用会导致超时。

修复：使用sys.stdin.read()一次性读取所有输入，再分割处理，大幅提升输入速度。

修复后的代码 #

import sys

def bellman_ford(n, edges):
    dist = [0] * (n + 1)  # 初始化所有节点距离为0，覆盖所有负环场景
    prev = [-1] * (n + 1)
    
    # 松弛阶段：加入提前终止优化
    for i in range(n):
        updated = False
        for u, v, w in edges:
            if dist[u] != float('inf') and dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
                prev[v] = u
                updated = True
        if not updated:
            break  # 无更新时提前退出，减少计算量
    
    # 检测负环并构建环路径
    for u, v, w in edges:
        if dist[u] != float('inf') and dist[u] + w < dist[v]:
            # 找到环中的一个节点
            current = v
            for _ in range(n):
                current = prev[current]
            
            # 回溯生成完整环
            cycle = []
            start = current
            while True:
                cycle.append(current)
                current = prev[current]
                if current == start:
                    cycle.append(current)
                    break
            return True, cycle
    
    return False, None

def main():
    data = sys.stdin.read().split()
    ptr = 0
    n = int(data[ptr])
    ptr +=1
    e = int(data[ptr])
    ptr +=1
    edges = []
    for _ in range(e):
        u = int(data[ptr])
        ptr +=1
        v = int(data[ptr])
        ptr +=1
        w = int(data[ptr])
        ptr +=1
        edges.append((u, v, w))
    
    chk, cycle = bellman_ford(n, edges)
    if chk:
        print("YES")
        print(*cycle)
    else:
        print("NO")

if __name__ == "__main__":
    main()

额外优化说明 #

• 松弛阶段加入updated标记，当某一轮无距离更新时提前退出，减少不必要的计算。
• 环的查找逻辑更健壮，确保能正确提取完整的负环路径。

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