多测试用例下2-SAT代码仅输出首个结果的问题排查

问题背景 #

我在做一道求解布尔方程的题目（属于2-SAT问题），输入包含多个测试用例，用sys.stdin.readlines()读取全部输入。单独处理单个测试用例时代码正常，但多测试用例时只输出第一个结果，后面的都被忽略。

原代码 #

import sys

def dfs(word, some_list):
    if len(word) == n + 1:
        sat(word, some_list)
        return
    for a in range(2):
        dfs(word+[a], some_list)


def sat(word, some_list):
    for xi, xj in some_list:
        element_1 = word[xi] if xi > 0 else not word[-xi]
        element_2 = word[xj] if xj > 0 else not word[-xj]
        if (element_1 or element_2) == 0:
            return
    for i in range(1, n+1):
        print(word[i], end='')
    exit()


data = sys.stdin.readlines()
id = 0

while id < len(data) - 1:
    n, m = map(int, data[id ].split(" "))
    id += 1
    some_list = []
    for _ in range(m):
        x, y = map(int, data[id ].split(" "))
        some_list .append([x, y])
        id += 1)
    dfs([], some_list)

测试用例与输出情况 #

测试输入 #

3 4 #First test
1 -2
-1 3
1 -3
-1 1
3 4 #Second test
-1 2
-2 3
1 3
3 2
1 0 #Third test

预期输出 #

000 #First test result
001 #Second test result
0   #Third test result

实际输出 #

000 #First test result

问题原因与解决方法 #

问题核心是sat函数里的exit()调用——找到第一个测试用例的可行解后，程序直接退出，后续测试用例完全没机会执行。

修改后的代码 #

import sys

def dfs(word, some_list):
    if len(word) == n + 1:
        if sat(word, some_list):
            return True
        return False
    for a in range(2):
        if dfs(word + [a], some_list):
            return True
    return False


def sat(word, some_list):
    for xi, xj in some_list:
        element_1 = word[xi] if xi > 0 else not word[-xi]
        element_2 = word[xj] if xj > 0 else not word[-xj]
        if not (element_1 or element_2):
            return False
    for i in range(1, n+1):
        print(word[i], end='')
    print()  # 换行分隔不同测试用例结果
    return True


data = sys.stdin.readlines()
id = 0

# 修正循环条件，避免漏掉最后一个测试用例
while id < len(data):
    # 跳过输入中的空行，提升鲁棒性
    while id < len(data) and data[id].strip() == '':
        id += 1
    if id >= len(data):
        break
    n, m = map(int, data[id].split())
    id += 1
    some_list = []
    for _ in range(m):
        while id < len(data) and data[id].strip() == '':
            id += 1
        x, y = map(int, data[id].split())
        some_list.append([x, y])
        id += 1
    dfs([], some_list)

修改说明 #

1. 移除exit()，改为让sat函数返回True标记找到可行解，False表示当前赋值不满足条件
2. 修改dfs函数，一旦子递归返回True（找到解），立刻向上返回终止递归，不再尝试其他分支
3. 修正主循环条件为id < len(data)，避免漏掉最后一个测试用例
4. 添加空行处理逻辑，应对输入中可能存在的空行情况
5. 打印结果后添加换行，让不同测试用例的输出更清晰

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