[python刷题模板] 数位DP
一、 算法&数据结构
1. 描述
数位DP我一直很头疼,直到用了灵神的记忆化模板。
-
用记忆化搜索的话,复杂度比数组形式dp高一些,但是胜在思路清晰,套路固定。
-
由于复杂度取决于状态数,在int下,位数最多就是10,所以重头应该在mask,一般都能过。
def f(i,mask,is_limit,is_num):
- i代表当前枚举到第几位
- mask可能需要也可能不需要,代表题目特有的限制条件。
- is_limit代表当前位能填的数是否受到了前一位的限制,这里通常是高位限制;但有的题也可能有低位限制。比如lc1742,下边有写。
- is_num代表前边是否填了数字,有的题目不需要这个参数,但通常可以写上。当i==n时,填了位代表是一组合法条件,返回1,;否则返回0。
2. 复杂度分析
- 取决于状态数。
3. 常见应用
- 计算方案数。
- 累计数字满足的属性数目,比如所有范围内数字中1的数目累计。
4. 常用优化
- 如果是高低位都限制的题目,可以把小的数字位数向高位补齐。
二、 模板代码
1. 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
2. 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)
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))
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)
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)
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)
三、其他
- 如果还是卡常数,有些区间问题可以转化为树状数组,常数小,代码短,不过真的很难理解,还是线段树好写。遇到就套板吧。