[python刷题模板] DFS回溯枚举

一、 算法&数据结构

1. 描述

dfs枚举是经常会用到的算法，比如小球放盒子，枚举所有组合情况。

2. 复杂度分析

1. O(xⁿ)
2. O(n!)
3. O(C(n,m))

3. 常见应用

1. 指数型枚举
2. 排列型枚举
3. 组合型枚举

4. 常用优化

在这之前，请千万不要忘记我们用的是python！
指数：

class Solution(object):
    def letterCasePermutation(self, S):
        f = lambda x: (x.lower(), x.upper()) if x.isalpha() else x
        # print(list(map(f,S)))
        # print(*map(f,S))
        return list(map("".join, itertools.product(*map(f, S))))

排列：

class Solution:
    def permute(self, nums: List[int]) -> List[List[int]]:       
        return list(itertools.permutations(nums) )

组合：

class Solution:
    def combine(self, n: int, k: int) -> List[List[int]]:        
        return list(itertools.combinations(range(1,n+1),k))

1. 剪枝

二、 模板代码

1. 三种枚举

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File       :   DFS.py    
@License    :   (C)Copyright 2013-2022, capstone
@Project    :   NormalTools
@Software   :   PyCharm
@ModifyTime :   2022/5/31 17:12    
@Author     :   liushuliang
@Version    :   1.0
@Description:   None
"""

def dfs_exp(step, arr, result, mask=True):
    """
    指数型枚举，每个元素都可以存在或不存在
    枚举每个元素
    :param step: 枚举第几个数
    :param arr: 原数组
    :param mask: 剪枝条件，用mask判断这个数是否需要放
    :return:
    """
    if step == len(arr):
        print(result)
        return
    dfs_exp(step + 1, arr, result[:], mask)  # 不选本数，直接处理下一步
    new_mask = mask and True  # 通过mask判断是否选本数
    if new_mask:
        dfs_exp(step + 1, arr, result + [a[step]], new_mask)  # 如果选本数，mask更新


a = [1, 2, 3]
n = len(a)


# dfs_exp(0,a,[],True)

def dfs_permutation(step, arr, result, help, mask=True):
    """
    排列型枚举，每个元素都要存在，枚举所有位置
    枚举每个位置，放每个数，回溯
    :param step: 枚举第几个位置
    :param arr: 原数组
    :param mask: 剪枝条件，用mask判断这个数是否需要放
    :return:
    """
    if step == len(arr):
        print(result)
        return

    for i in range(len(arr)):
        if help[i] == 0:
            help[i] = 1
            dfs_permutation(step + 1, arr, result + [arr[i]], help, mask)
            help[i] = 0


a = [1, 2, 3, 4]
help_arr = [0] * n


# dfs_permutation(0,a,[],help_arr)

def dfs_combination(step, start, arr, limit, result, mask=True):
    """
    组合型枚举,从arr中任选limit个数，顺序不重要
    :param step:
    :param arr:
    :param result:
    :param help:
    :param mask:
    :return:
    """
    if step == limit:
        print(result)
        return
    for i in range(start, len(arr)):
        dfs_combination(step + 1, i + 1, arr, limit, result + [arr[i]])


a = [ 2, 3, 4,1]
dfs_combination(0, 0, a, 3, [])

2.指数型枚举

链接: 473. 火柴拼正方形
链接: [LeetCode解题报告] 473. 火柴拼正方形

class Solution:
    def makesquare(self, matchsticks: List[int]) -> bool:
        n = len(matchsticks)
        if n < 4:
            return False
        s = sum(matchsticks)
        side = s // 4
        if side * 4 != s:
            return False
        matchsticks.sort(reverse=True)
        res = [0] * 4
        # 枚举每个数字是否出现能组合成什么值,2^n;这个组合值用于下边dfs的剪枝：
        # 如果放火柴之后，本边剩余长度不能用火柴组合起来，那直接剪掉
        # 本来我就是试试，没想到剪成功了！
        cmb = set()
        def dfs_sum(i,s):
            if i == n:
                cmb.add(s)
                return 
            dfs_sum(i+1,s)
            dfs_sum(i+1,s+matchsticks[i])
        dfs_sum(0,0)
        # print(cmb)
        def dfs(i):
            if i == n:
                return True
                # return all(res[j]==side for j in range(4))
            
            for j in range(4):
                if  res[j]+matchsticks[i]>  side or side - res[j]-matchsticks[i] not in cmb :                    
                    continue
                res[j] += matchsticks[i]               
                if dfs(i+1):
                    return True
                res[j] -= matchsticks[i]
            return False

        return dfs(0)

784. 字母大小写全排列

class Solution:
    def letterCasePermutation(self, s: str) -> List[str]:
        n = len(s)
        ans = []
        ret = ['']*n
        def dfs(pos):
            if pos ==n :
                return ans.append(''.join(ret))
            c = s[pos]
            if c.isdigit():
                ret[pos] = c 
                dfs(pos + 1)
            else:
                ret[pos] = s[pos].lower()
                dfs(pos+1)
                ret[pos] = s[pos].upper()
                dfs(pos+1)
        dfs(0)
        return ans

3. 排列型枚举

简单点
46. 全排列

class Solution:
    def permute(self, nums: List[int]) -> List[List[int]]:
        n = len(nums)
        used = [0]*n
        ans = []
        def dfs(s):
            if len(s) == n:
               ans.append(s[:])
               return 
            for i in range(n):
                if not used[i]:
                    used[i] = 1
                    dfs(s+nums[i:i+1])
                    used[i] = 0
        dfs([])
        return ans 

    a = [1,2,3,4]
    n = len(a)
    used = [0]*n
    ret = [0]*n
    ans = []

    def dfs(i):
        if i == n:
            ans.append(ret[:])
            # print(ret)

            return
        for j in range(n):
            if used[j] == 0:
                used[j] = 1
                ret[i] = a[j]
                dfs(i+1)
                used[j] = 0

    dfs(0)

4.组合型枚举

77. 组合

class Solution:
    def combine(self, n: int, k: int) -> List[List[int]]:        
        ans = []
        def dfs(start,s):
            if len(s) == k:
                ans.append(s[:])
                return 
            for i in range(start,n):                
                dfs(i+1,s+[i+1])
        
        dfs(0,[])
        return ans 

三、其他

四、更多例题

五、参考链接