[python刷题模板] 数位DP

一、 算法&数据结构

1. 描述

数位DP我一直很头疼，直到用了灵神的记忆化模板。
1

• 用记忆化搜索的话，复杂度比数组形式dp高一些，但是胜在思路清晰，套路固定。

• 由于复杂度取决于状态数，在int下，位数最多就是10，所以重头应该在mask，一般都能过。

  def f(i,mask,is_limit,is_num):

  • i代表当前枚举到第几位
  • mask可能需要也可能不需要，代表题目特有的限制条件。
  • is_limit代表当前位能填的数是否受到了前一位的限制，这里通常是高位限制；但有的题也可能有低位限制。比如lc1742,下边有写。
  • is_num代表前边是否填了数字，有的题目不需要这个参数，但通常可以写上。当i==n时，填了位代表是一组合法条件，返回1,；否则返回0。

2. 复杂度分析

1. 取决于状态数。

3. 常见应用

1. 计算方案数。
2. 累计数字满足的属性数目，比如所有范围内数字中1的数目累计。

4. 常用优化

1. 如果是高低位都限制的题目，可以把小的数字位数向高位补齐。

二、 模板代码

1. 357. 统计各位数字都不同的数字个数

例题: 357. 统计各位数字都不同的数字个数

• 这题有数字次数限制和高位限制。
• 其中mask是一个10位的状压，代表0~9的数字是否已使用。枚举当前位数字j累加ans时只扫描未使用的位。
• 由于高位不允许填0，因此要特殊处理前边位没填（即is_num=False），且本位也不填的情况。

class Solution:
    def countNumbersWithUniqueDigits(self, n: int) -> int:
        if n == 0:
            return 1
        s = str(10**n-1)
        n = len(s)
        @cache
        def f(i,mask,is_limit,is_num):
            if i == n:
                return int(is_num)
            up = int(s[i]) if is_limit else 9
            down = 0 if is_num else 1
            ans = 0
            if not is_num:
                ans += f(i+1,0,False,False)
            for j in range(down,up+1):
                if mask & (1<<j) == 0:
                    ans  += f(i+1,mask|(1<<j),is_limit and j == up,True)
            return ans 
        return f(0,0,True,False)+1
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

2. 1012. 至少有 1 位重复的数字

链接: 1012. 至少有 1 位重复的数字

由于至少有一位代表可以有多位，不好分析，于是直接转化成没有重复，变成上一道题。

class Solution:
    def numDupDigitsAtMostN(self, n: int) -> int:
        x = n
        s = str(n)
        n = len(s)
        @cache
        def f(i,mask,is_limit,is_num):
            if i == n:
                return int(is_num)
            up = int(s[i]) if is_limit else 9
            down = 0 if is_num else 1
            ans = 0
            if not is_num:
                ans += f(i+1,0,False,False)
            for j in range(down,up+1):
                if mask & (1<<j) == 0:
                    ans  += f(i+1,mask|(1<<j),is_limit and j == up,True)
            return ans 
        # print(x,f(0,0,True,False))
        return x- f(0,0,True,False)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

3. 同时限制上下界的数位DP。

链接: 1742. 盒子中小球的最大数量
参见: [LeetCode解题报告] 1742. 盒子中小球的最大数量](https://blog.csdn.net/liuliangcan/article/details/127994903)

• 由于数据小，可以模拟过。这里写数位DP，可以支持更大的数据范围。
• 题目实际上是求各种数位和对应的数字数量，取max。由于数位和就是1-46，因此这个值可以作为mask传进去。
• 退出条件要注意，如果mask<0代表非法，即用了前边位，后边位的和要求是负数，直接返回0即可；同时如果到n，要判断和是否用完了，即mask==0。

class Solution:
    def countBalls(self, lowLimit: int, highLimit: int) -> int:
        s1 = str(highLimit)
        s2 = str(lowLimit)
        
        n = len(s1)
        if len(s2)< n:  # 前边补齐0
            s2 = '0'*(n-len(s2)) + s2
        
        @cache
        def f(i,s,is_up_limit,is_down_limit,is_num):  # 枚举到第几位，从i往后的位和，上限是否受到前一位限制，下限是否受到前一位限制，前边是否填过数;返回这个s下的方案数
            if s < 0:
                return 0
            if i == n:
                return int(is_num and s == 0)
            up = int(s1[i]) if is_up_limit else 9 
            down = int(s2[i]) if not is_num or is_down_limit else 0  # 前一位没填，或者下限受限，down都=s2
            ans = 0
            for j in range(down,up+1):
                ans += f(i+1,s-j,j == up and is_up_limit,is_down_limit and j == down,is_num or j>0 )            
            return ans
        # for i in range(46):
        #     print(i,f(0,i,True,False))
        # print(f(0,6,True,True,False))
        return max(f(0,i,True,True,False) for i in range(46))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

4. 233. 数字 1 的个数

链接: 233. 数字 1 的个数
链接: [面试题 17.06. 2出现的次数](面试题 17.06. 2出现的次数)

• mask变成计数前缀1的数量，最后返回。

class Solution:
    def countDigitOne(self, x: int) -> int:
        s = str(x)
        n = len(s)
        @cache
        def f(i,cnt1,is_limit,is_num):
            if i == n:
                return cnt1
            ans = 0
            # if not is_num:
            #     ans += f(i+1,cnt1,False,False)
            r = int(s[i]) if is_limit else 9
            # l = 0 if is_num else 1
            l = 0
            for j in range(l,r+1):                
                ans += f(i+1,cnt1+(j==1),is_limit and j == r,True)
            return ans 
        return f(0,0,True,False)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

5. 合法数字本身受限 600. 不含连续1的非负整数

链接: 600. 不含连续1的非负整数

• 数字转化成2进制。
• mask变成前一位是否是1，如果是1后一位不能填，否则后一位可以填。
• 由于二进制只有01就不循环枚举j了，注意只有当前一位不是1且当前位上线是1才能填1.

class Solution:
    def countDigitOne(self, x: int) -> int:
        s = str(x)
        n = len(s)
        @cache
        def f(i,cnt1,is_limit,is_num):
            if i == n:
                return cnt1
            ans = 0
            # if not is_num:
            #     ans += f(i+1,cnt1,False,False)
            r = int(s[i]) if is_limit else 9
            # l = 0 if is_num else 1
            l = 0
            for j in range(l,r+1):                
                ans += f(i+1,cnt1+(j==1),is_limit and j == r,True)
            return ans 
        return f(0,0,True,False)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

6. 可用的数位数字受限 902. 最大为 N 的数字组合

链接: 902. 最大为 N 的数字组合

• 数字不是0-9了，而是给出了只能使用这些。
• 类似。

class Solution:
    def atMostNGivenDigitSet(self, digits: List[str], x: int) -> int:
        s = str(x)
        n = len(s)
        @cache
        def f(i,is_limit,is_num):
            if i == len(s):
                return int(is_num)
            ans = 0
            if not is_num:
                ans += f(i+1,False,False)
            up = s[i] if is_limit else '9'
            for j in digits:
                if j <= up:
                    ans += f(i+1,is_limit and j == up,True)
            return ans
            
          
        return f(0,True,False)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

三、其他

1. 如果还是卡常数，有些区间问题可以转化为树状数组，常数小，代码短，不过真的很难理解，还是线段树好写。遇到就套板吧。

四、更多例题

• [LeetCode解题报告] 2376. 统计特殊整数

五、参考链接