【算法】回溯算法专题③ ——排列型回溯 python

news2025/2/4 20:07:09

前置

【算法】回溯算法专题① ——子集型回溯 python
【算法】回溯算法专题② ——组合型回溯 + 剪枝 python

小试牛刀

全排列 https://leetcode.cn/problems/permutations/description/

给定一个不含重复数字的数组 nums ，返回其所有可能的全排列。你可以按任意顺序返回答案。

示例 1：

输入：nums = [1,2,3]
输出：[[1,2,3],[1,3,2],[2,1,3],[2,3,1],[3,1,2],[3,2,1]]

示例 2：

输入：nums = [0,1]
输出：[[0,1],[1,0]]
示例 3：

输入：nums = [1]
输出：[[1]]

提示：

1 <= nums.length <= 6
-10 <= nums[i] <= 10
nums 中的所有整数互不相同

思路:

用一个vis标记数组:对已选的数字打上标记

题解:

class Solution:
    def permute(self, nums: List[int]) -> List[List[int]]:
        n = len(nums)
        path = []
        ans = []
        vis = [0] * (n + 1)

        def dfs(i):
            if i == n:
                ans.append(path.copy())
                return
            for j in range(n):
                if not vis[j]:
                    vis[j] = 1
                    path.append(nums[j])
                    dfs(i + 1)
                    path.pop()
                    vis[j] = 0 # 别忘了回溯时将vis打回0

        dfs(0)
        return ans

当然vis标记数组可以通过递归时传入一个set参数代替

题解2:

class Solution:
    def permute(self, nums: List[int]) -> List[List[int]]:
        n = len(nums)
        path = []
        ans = []

        def dfs(i, s): # s:set
            if i == n:
                ans.append(path.copy())
                return
            for x in s:
                path.append(x)
                dfs(i + 1, s - {x})
                path.pop()

        dfs(0, set(nums)) # 用集合可以 O(1)判断元素是否在里面
        return ans

回归经典

N皇后 https://leetcode.cn/problems/n-queens/description/

按照国际象棋的规则，皇后可以攻击与之处在同一行或同一列或同一斜线上的棋子。

n 皇后问题研究的是如何将 n 个皇后放置在 n×n 的棋盘上，并且使皇后彼此之间不能相互攻击。

给你一个整数 n ，返回所有不同的 n 皇后问题的解决方案。

每一种解法包含一个不同的 n 皇后问题的棋子放置方案，该方案中 ‘Q’ 和 ‘.’ 分别代表了皇后和空位。

示例 1：

在这里插入图片描述
输入：n = 4
输出：[[“.Q…”,“…Q”,“Q…”,“…Q.”],[“…Q.”,“Q…”,“…Q”,“.Q…”]]
解释：如上图所示，4 皇后问题存在两个不同的解法。

示例 2：

输入：n = 1
输出：[[“Q”]]

提示：
1 <= n <= 9

思路:

同一对角线不能有多个皇后
（可以通过计算行号加或减列号来判断）

在这里插入图片描述

题解code:

class Solution:
    def solveNQueens(self, n: int) -> List[List[str]]:
        ans = []
        col = [0] * n  # col:列
        vis = [0] * n  # vis:标记数组
        diag1 = [0] * (2 * n)
        diag2 = [0] * (2 * n)
        # 分别表示两个对角线

        def dfs(r):  # row:行
            if r == n:
                ans.append(["." * c + "Q" + "." * (n - 1 - c) for c in col])
                return
            for c in range(n):
                if not vis[c] and not diag1[r + c] and not diag2[r - c]:
                    col[r] = c
                    vis[c] = diag1[r + c] = diag2[r - c] = 1
                    dfs(r + 1)
                    vis[c] = diag1[r + c] = diag2[r - c] = 0

        dfs(0)
        return ans

举一反三

小小变化一下,可以在对角线上,但有前提: 行号至少相差3

受伤的皇后 https://www.lanqiao.cn/problems/552/learning/

在这里插入图片描述

思路:

主要问题在于判断相同对角线上行号之差
以(2, 2)为例：
【右对角线】的点是 (1, 3)，【左对角线】的点是 (3, 3)，
可以发现：
(2, 2)和 (1, 3)的横坐标之差的绝对值=纵坐标之差的绝对值
同样的(2, 2)和 (3, 3)亦是如此，
所以判断两点是否在同一对角线，
即判断这两点的横纵坐标绝对值之差是否相等

题解code:

col = [0] * n
vis = [0] * n
diag1 = [0] * 2 * n
diag2 = [0] * 2 * n

def check(x, y):
    if not vis[y] and not diag1[x + y] and not diag2[x - y]:
        for i in range(x):
            if abs(i - x) == abs(col[i] - y) and abs(x - i) < 3:
                return 0
        return 1
    return 0

def dfs(r):
    if r == n:
        global ans
        ans += 1
        return
    for c in range(n):
        if check(r, c):
            col[r] = c
            vis[c] = diag1[r + c] = diag2[r - c] = 1
            dfs(r + 1)
            col[r] = 0
            vis[c] = diag1[r + c] = diag2[r - c] = 0
            
n = int(input())
ans = 0
dfs(0)
print(ans)