回溯算法——N皇后问题

一、背景介绍

根据国际象棋的规则，皇后可以攻击与之处在同一行或同一列或同一斜线上的棋子。给定𝑛 个皇后和一个𝑛×𝑛大小的棋盘，寻找使得所有皇后之间无法相互攻击的摆放方案。

如图所示，当𝑛=4时，共可以找到两个解。从回溯算法的角度看，𝑛×𝑛大小的棋盘共有个格子，给出了所有的选择 choices 。在逐个放置皇后的过程中，棋盘状态在不断地变化，每个时刻的棋盘就是状态state。

下图展示了本题的三个约束条件：多个皇后不能在同一行、同一列、同一对角线。值得注意的是，对角线分为主对角线\和次对角线/两种。

1. 逐行放置策略

皇后的数量和棋盘的行数都为𝑛，因此我们容易得到一个推论：棋盘每行都允许且只允许放置一个皇后。

也就是说，我们可以采取逐行放置策略：从第一行开始，在每行放置一个皇后，直至最后一行结束。

如图所示，为4皇后问题的逐行放置过程。受画幅限制，下图仅展开了第一行的其中一个搜索分 支，并且将不满足列约束和对角线约束的方案都进行了剪枝。

本质上看，逐行放置策略起到了剪枝的作用，它避免了同一行出现多个皇后的所有搜索分支。

2. 列与对角线剪枝

为了满足列约束，我们可以利用一个长度为𝑛的布尔型数组cols 记录每一列是否有皇后。 在每次决定放置前，我们通过 cols 将已有皇后的列进行剪枝，并在回溯中动态更新 cols 的状态。

那么，如何处理对角线约束呢？设棋盘中某个格子的行列索引为(𝑟𝑜𝑤,𝑐𝑜𝑙)，选定矩阵中的某条主对角线， 我们发现该对角线上所有格子的行索引减列索引都相等，即对角线上所有格子的𝑟𝑜𝑤−𝑐𝑜𝑙为恒定值。 也就是说，如果两个格子满足𝑟𝑜𝑤1−𝑐𝑜𝑙1 =𝑟𝑜𝑤2−𝑐𝑜𝑙2，则它们一定处在同一条主对角线上。

利用该规律，我们可以借助下图所示的数组 diag1 ，记录每条主对角线上是否有皇后。 同理，次对角线上的所有格子的𝑟𝑜𝑤+𝑐𝑜𝑙是恒定值。我们同样也可以借助数组 diag2 来处理次对角线约束。

代码实现

请注意，𝑛维方阵中𝑟𝑜𝑤−𝑐𝑜𝑙的范围是[−𝑛+1,𝑛−1]，𝑟𝑜𝑤+𝑐𝑜𝑙的范围是[0,2𝑛−2]，所以主对 角线和次对角线的数量都为2𝑛−1，即数组 diag1 和 diag2 的长度都为2𝑛−1。

def backtrack(
    row: int,
    n: int,
    state: list[list[str]],
    res: list[list[list[str]]],
    cols: list[bool],
    diags1: list[bool],
    diags2: list[bool],
):
    """回溯算法：n 皇后"""
    # 当放置完所有行时，记录解
    if row == n:
        res.append([list(row) for row in state])
        return
    # 遍历所有列
    for col in range(n):
        # 计算该格子对应的主对角线和次对角线
        diag1 = row - col + n - 1
        diag2 = row + col
        # 剪枝：不允许该格子所在列、主对角线、次对角线上存在皇后
        if not cols[col] and not diags1[diag1] and not diags2[diag2]:
            # 尝试：将皇后放置在该格子
            state[row][col] = "Q"
            cols[col] = diags1[diag1] = diags2[diag2] = True
            # 放置下一行
            backtrack(row + 1, n, state, res, cols, diags1, diags2)
            # 回退：将该格子恢复为空位
            state[row][col] = "#"
            cols[col] = diags1[diag1] = diags2[diag2] = False


def n_queens(n: int) -> list[list[list[str]]]:
    """求解 n 皇后"""
    # 初始化 n*n 大小的棋盘，其中 'Q' 代表皇后，'#' 代表空位
    state = [["#" for _ in range(n)] for _ in range(n)]
    cols = [False] * n  # 记录列是否有皇后
    diags1 = [False] * (2 * n - 1)  # 记录主对角线上是否有皇后
    diags2 = [False] * (2 * n - 1)  # 记录次对角线上是否有皇后
    res = []
    backtrack(0, n, state, res, cols, diags1, diags2)

    return res


"""Driver Code"""
if __name__ == "__main__":
    n = 4
    res = n_queens(n)

    print(f"输入棋盘长宽为 {n}")
    print(f"皇后放置方案共有 {len(res)} 种")
    for state in res:
        print("--------------------")
        for row in state:
            print(row)

逐行放置𝑛次，考虑列约束，则从第一行到最后一行分别有𝑛、𝑛−1、…、2、1个选择，因此时间复杂度为𝑂(𝑛!)。实际上，根据对角线约束的剪枝也能够大幅地缩小搜索空间，因而搜索效率往往优于以上时间复杂度。

数组state使用𝑂()空间，数组 cols、 diags1 和 diags2 皆使用𝑂(𝑛)空间。最大递归深度为𝑛，使用𝑂(𝑛)栈帧空间。因此，空间复杂度为𝑂()。