如何使用Python求解满足行和列和约束的5x5二元矩阵

输入方式定义 #

你不需要输入带加号和未知量的完整矩阵，直接用两个列表分别存行和约束、列和约束即可：

• 行和约束列表：顺序对应从上到下每一行的目标求和值，示例里的约束写为row_sums = [1, 0, 1, 3, 2]
• 列和约束列表：顺序对应从左到右每一列的目标求和值，示例里的约束写为col_sums = [2, 1, 2, 1, 1]

如果你需要直接解析题目里那种带加号的文本格式矩阵，额外加个字符串解析逻辑提取这两个列表即可。

求解代码实现 #

我们不用暴力遍历所有2^25种可能，提前按行过滤符合行和的候选，再加列和剪枝，效率会高很多：

import itertools

def solve_binary_matrix(row_sums, col_sums):
    n = len(row_sums)
    # 预生成每行所有符合行和要求的0-1组合
    row_candidates = []
    for target_sum in row_sums:
        valid_rows = []
        for bits in itertools.product([0, 1], repeat=n):
            if sum(bits) == target_sum:
                valid_rows.append(bits)
        row_candidates.append(valid_rows)
    
    # 回溯搜索+剪枝
    def backtrack(row_index, current_col_sum, current_matrix):
        # 所有行遍历完成，校验列和是否完全匹配
        if row_index == n:
            return current_matrix if current_col_sum == col_sums else None
        # 遍历当前行所有候选
        for candidate_row in row_candidates[row_index]:
            new_col_sum = [current_col_sum[i] + candidate_row[i] for i in range(n)]
            # 剪枝：任意列累计和超过目标值就直接跳过，不用往下递归
            if any(new_col_sum[i] > col_sums[i] for i in range(n)):
                continue
            result = backtrack(row_index + 1, new_col_sum, current_matrix + [candidate_row])
            if result:
                return result
        return None

    return backtrack(0, [0]*n, [])

# 测试示例
if __name__ == "__main__":
    row_sums = [1, 0, 1, 3, 2]
    col_sums = [2, 1, 2, 1, 1]
    solution = solve_binary_matrix(row_sums, col_sums)
    if solution:
        print("找到可行解：")
        for i, row in enumerate(solution):
            print(f"{'+'.join(map(str, row))}| {row_sums[i]}")
        print("_"*15)
        print(" ".join(map(str, col_sums)))
    else:
        print("不存在符合约束的解")

运行结果示例 #

运行上述代码会输出符合要求的解，比如：

找到可行解：
0+0+0+0+1| 1
0+0+0+0+0| 0
0+0+0+1+0| 1
1+1+1+0+0| 3
1+0+1+0+0| 2
_______________
2 1 2 1 1

如果想要找到所有可行解，只需要把回溯逻辑里找到解就返回的逻辑改成把解存入全局列表，遍历完所有可能即可。

扩展说明 #

如果矩阵尺寸超过8x8，回溯的效率会明显下降，你可以改用pulp等线性规划库求解0-1规划问题，性能会提升几个数量级。

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