数位dp题集 Python

荷碧TongZJ

已于 2024-04-10 21:16:48 修改

阅读量808

点赞数 3

分类专栏：数据结构与算法文章标签： python 算法动态规划数位 dp

于 2023-04-02 00:07:35 首次发布

本文链接：https://blog.csdn.net/qq_55745968/article/details/129904606

版权

数据结构与算法专栏收录该内容

17 篇文章 207 订阅

订阅专栏

数位 dp 是个让人头疼的问题，特别是对于前导零的处理方法，还是得多多练习

数字计数 100🏆

https://www.luogu.com.cn/problem/P2602

【问题描述】

给定两个正整数 a 和 b，求在 [a, b] 中的所有整数中，每个数码各出现了多少次

【输入格式】

仅包含一行两个整数 a, b，含义如上所述

【输出格式】

包含一行十个整数，分别表示 0 ~ 9 在 [a,b] 中出现了多少次。

【样例】

输入	输出
1 99	9 20 20 20 20 20 20 20 20 20

【评测用例规模与约定】

30%	$a \leq b \leq 10^6$
100%	$1 \leq a \leq b \leq 10^{12}$

【解析及代码】

给定区间 [a, b] 求其中各个数出现的次数，可以等价于 [1, b] 中各个数出现的次数 - 区间 [1, a-1] 中各个数出现的次数

例如给定 b = 465，那么各个数出现的次数可以这样拆分来计算 (先不考虑前导零)：

百位 4：当百位为 0~3 时，000~399 的贡献
十位 6：固定百位为 4，当十位为 0~5 时，400~459 的贡献
个位 5：固定十位为 6，当个位为 0~5 时，460~465 的贡献

以百位 4 为例，将数值的贡献拆分为三种类型的贡献：

本位 ∈ [0, 4) 时，下位的贡献：百位以下的数位的贡献，即 000~399 中 4 个 00~99 的贡献
本位 ∈ [0, 4) 时，本位的贡献：百位不为最大值时的贡献，即 000~399 中的百位，0~3 这三个数字分别出现了 100 次
本位 = 4 时，本位的贡献：百位为最大值时的贡献，即 400 ~ 465 中的百位，4 出现了 65 次 (这个很关键)

所以，在对百位的贡献进行上述的拆分、计算后，可以忽略百位对十位、个位的影响

此外，在这个计算的过程中，又需要使用 0~9、00~99、000~999 …… 中各个数字出现的次数。因为先不考虑前导零 (所有数字一视同仁)，所以用一个一维列表即可完成动态规划

使用 dp[i] 表示 $[0, 10^i)$ 中数字 9 的出现次数，则 dp[3] 的贡献也分为两种类型 (dp[2] 表示 00~99 中数字 9 的出现次数)：

下位的贡献：由 00~99 新增加一位到 000~999，则新增加了 10 个 00~99 的贡献
本位的贡献：新增加的一位是 0~9，则 0~9 各出现了 100 次

可以得出以下状态转移方程：

$dp[i+1] = 10 \cdot dp[i] + 10^i$

最后就是删除前导零，比如 001 中的两个 0，029 中的一个 0。不难得出，百位贡献了 100 个前导零 (000~099)，十位贡献了 10 个前导零 (00~09)，个位贡献了 1 个前导零 (0)

a, b = input().split()

dp = [0] * len(b)
# dp[i] 表示 0 ~ 10^i-1 中 9 出现的次数
for i in range(len(dp) - 1):
    dp[i + 1] = dp[i] * 10 + 10 ** i


def solve(n):
    tmp = int(n) + 1
    cnt = [0] * 10
    if tmp > 1:
        # 从高位开始枚举: i 为本位的位置
        for i, digit in zip(range(len(n) - 1, -1, -1), map(int, n)):
            gain = 10 ** i
            # 下位的贡献: 当本位 ∈ [0, digit)
            for j in range(10): cnt[j] += dp[i] * digit
            # 本位的贡献: 当本位 ∈ [0, digit)
            for j in range(digit): cnt[j] += gain
            # 本位的贡献: 当本位 = digit
            tmp -= digit * gain
            cnt[digit] += tmp
            # 删除前导零
            cnt[0] -= gain
    return cnt


print(*[j - i for i, j in zip(solve(str(int(a) - 1)), solve(b))])

Windy 数 100🏆

https://loj.ac/p/10165

【问题描述】

Windy 定义了一种 Windy 数：不含前导零且相邻两个数字之差至少为 2 的正整数被称为 Windy 数。

Windy 想知道，在 A 和 B 之间，包括 A 和 B，总共有多少个 Windy 数？

【输入格式】

一行两个数，分别为 A，B。

【输出格式】

输出一个整数，表示答案。

【样例】

输入	输出
1 10	9
25 50	20

【评测用例规模与约定】

20%	$1 \leq A \leq B \leq 10^6$
100%	$1 \leq A \leq B \leq 2 \times 10^9$

【解析及代码】

把问题做一下等价：

区间 [A, B] 之间的 windy 数等价于：[1, B] 内的 windy 数减去 [1, A] 内的 windy 数
如果给定 B = 9654，那么可以拆分为几部分贡献：0~8999、9000~9599、9600~9649、9650~9654。但显而易见的是，9650~9654 因为有 6 和 5 相邻，所以这部分没有 windy 数

那么给定 B = 96544，应该对这个数字做以下等价：

千位 6 - 2 < 百位 5 < 千位 6 + 2：将百位 5 改为 4 (千位 - 2)，同时因为百位减少，所以后续的数位应改为最大值，即 B 等价于 96499
十位 9 - 2 < 个位 9 < 十位 9 + 2：将个位 9 改为 7 (十位 - 2)，即 B 等价于 96497

然后，定义二维矩阵 dp，用 dp[i][j] 表示第 i 个数位 (0 为个位) 为 j 时的 windy 数的个数，状态转移方程为：

$dp[i][j] = \sum_{k=0}^{j - 2}dp[i-1][k] + \sum_{k=j+2}^{10}dp[i-1][k]$

然后，还要注意这个等价方法需要借位。e.g.，1234 等价于 979

试着求解 [1, B] 内的 windy 数：

0~89999： $\sum_{j=1}^8 dp[4][j]$ 表示五位数、在 9 万以内的 windy 数， $\sum_{i=0}^3 \sum_{j=1}^9 dp[i][j]$ 表示 1 万以内的 windy 数，该区间内的 windy 数的个数为 $\sum_{j=1}^8 dp[4][j]+\sum_{i=0}^3 \sum_{j=1}^9 dp[i][j]$
90000~95999：固定万位为 9，该区间内的 windy 数为 $\sum_{j=0}^5 dp[3][j]$ ，即千位为 0~5 时的 windy 数个数
96000~96399：固定千位为 6，该区间内的 windy 数为 $\sum_{j=0}^3 dp[2][j]$ ，即百位为 0~3 时的 windy 数个数

以此类推，便可求解出答案

def equivalent(n):
    n = list(map(int, str(n)))
    # 进行等价处理: 96544 -> 96497
    for i in range(1, len(n)):
        if abs(n[i - 1] - n[i]) < 2:
            # 上一个数位为 0 / 1, 借位
            if n[i - 1] == 0:
                n[i - 2:i + 1] = n[i - 2] - 1, 9, 7
            elif n[i - 1] == 1:
                n[i - 1:i + 1] = 0, 9
            else:
                n[i] = n[i - 1] - 2
            # 修改后续数位的上限
            for j in range(i + 1, len(n)): n[j] = 9
    return n if n[0] else n[1:]


a, b = map(int, input().split())
a, b = map(equivalent, (a - 1, b))
# dp[0] 表示只有个位时的 windy 数
dp = [[0] * 10 for _ in range(len(b))]
dp[0] = [1] * 10


def search(i, digit, uplim=10):
    ''' i: 本位所在的数位
        digit: 本位的数字
        uplim: 下位的数字限制'''
    dolim = min(uplim, max(0, digit - 1))
    return sum(dp[i - 1][:dolim]) + sum(dp[i - 1][digit + 2:uplim])


# 枚举数位: 从低位到高位
for i in range(1, len(b)):
    # 枚举本位的数字
    for j in range(10):
        dp[i][j] = search(i, j)


def solve(n):
    length = len(n)
    if length <= 1: return n[0] if n else 0
    n.reverse()
    # 对于 464: 拆分为 0~399, 400~459, 460~464
    # 预计算 0~399 的部分
    ret = sum(dp[length - 1][1:n[-1]]) + sum(sum(dp[i][1:]) for i in range(length - 1))
    # 限定最高位为最大值, 计算其它部分
    for i in range(length - 1, 0, -1):
        if n[i - 1] < 0: break
        ret += search(i, n[i], uplim=n[i - 1] + (i == 1))
    return ret


print(solve(b) - solve(a))

数字1的个数 100🏆

https://leetcode.cn/problems/number-of-digit-one/

【问题描述】

给定一个整数 n，计算所有小于等于 n 的非负整数中数字 1 出现的个数。

【样例】

输入	输出
n = 13	6
n = 0	0

【评测用例规模与约定】

100%	$0 \leq n \leq 10^9$

【解析及代码】

定义一维矩阵 dp，以 dp[i] 表示 $[0, 10^i)$ 内数字 1 出现的个数，状态转移方程为：

$dp[i] = 10 \cdot dp[i-1] + 10^{i-1}$

以 n = 3125 为例，数字 1 出现的个数可划分为：

0 ~ 2999：千位以下 (下位) 的贡献为 $3 \cdot dp[3]$ ，千位 (本位) 的贡献为 $10^3$ (1000 ~ 1999 中千位的 1)
3000 ~ 3099：百位以下 (下位) 的贡献为 $1 \cdot dp[2]$ ，百位 (本位) 的贡献为 0
3100 ~ 3119：十位以下 (下位) 的贡献为 $2 \cdot dp[1]$ ，十位 (本位) 的贡献为 $10^1$ ；修正百位的贡献为 20
3120 ~ 3125：个位 (本位) 的贡献为 $10^0$ ；修正百位的贡献为 26

class Solution(object):
    def countDigitOne(self, n):
        n = str(n)
        length = len(n)
        # dp[i] 对应 [0, 10^i) 内数字 1 的个数
        dp = [0] * length
        for i in range(1, length):
            dp[i] = dp[i - 1] * 10 + 10 ** (i - 1)
        # 从高位到低位逐渐固定数字, 进行求解
        ret = 0
        for i in range(length, 0, -1):
            digit = int(n[length - i])
            # 下位的贡献
            ret += dp[i - 1] * digit
            # 本位的贡献
            if digit >= 1:
                ret += 10 ** (i - 1) if digit > 1 else \
                    (int(n[length - i + 1:] if i > 1 else 0) + 1)
        return ret