「数位dp」统计整数数目(力扣第2719题)

本题为1月16日力扣每日一题

题目来源:力扣第2719题

题目tag:数位dp 动态规划

题面

题目描述

给你两个数字字符串num1和num2，以及两个整数max_sum和min_sum。如果一个整数x满足以下条件，我们称它是一个好整数：

\(num1 \leq x \leq num2\)

\(min\_sum \leq digit\_sum(x) \leq max\_sum\)

请你返回好整数的数目。答案可能很大，请返回答案对$ 10^9 + 7 $取余后的结果。

注意，digit_sum(x)表示x各位数字之和。

示例

示例 1

输入：

num1 = "1", num2 = "12", min_num = 1, max_num = 8

输出：

11

解释:

总共有11个整数的数位和在1到8之间，分别是1,2,3,4,5,6,7,8,10,11和12。所以我们返回11。

示例 2

输入：

num1 = "1", num2 = "5", min_num = 1, max_num = 5

输出：

5

解释:

数位和在1到5之间的5个整数分别为1,2,3,4和5。所以我们返回5。

提示

\(1 \leq num1 \leq num2 \leq 10^{22}\)
\(1 \leq min\_sum \leq max\_sum \leq 400\)

---

思路分析

如果你此前有做过数位dp相关的题目的话,此题一眼就知道应该使用数位dp来完成.此题满足了所有数位dp适用的条件(整数较大,有一个上界和一个正数下界,所求需要数位满足一些性质),所以可以用数位dp来完成.

这里对所求数的范围限制(即 \(num1 \leq x \leq num2\) )可以完全按照数位dp的常规处理方法,用前缀和``差分思想拆成满足 \(x \leq num2\) 的x个数减去满足 \(x \leq num1 - 1\) 的x个数.由于num1为字符串,所以此处的 \(-1\) 操作可用高精度减法来完成.

数位dp有很多种写法(模板,或者叫公式),可以按照个人喜好套用自己擅长的写法,这里我简单介绍一下我自己的写法.

我个人的数位dp分成两个部分,且相对独立,分别如下:

第一部分,给定一个确定的上界,讨论每一位有几种取法.可能你会有疑惑,对于每一位,显然它可以取0到上界这一位的所有数啊?对但不完全对.可以注意到,如果这一位取0到上界这一位-1的数,那么这一位后面的每一位,它都可以任意选择,不受上界限制,此时无需讨论后面的位,可以直接得出答案(答案就是第二部分所求的结果).如果这一位取的是上界这一位的数,那么后续位仍然受到限制,此时需要继续讨论后面的位.因此我们可以发现,对数位的讨论其实是一颗二叉树,左叉是取0到上界这一位-1,这个情况无需继续分叉,可以直接求解;而右叉则是取上界这一位,此时需要继续分叉,不断分小情况求解.这一部分在不同的题目中几乎无需变化,可以完全按一样的写法写.

第二部分,对前面所提到的二叉树的左叉进行直接求解.如何求解?当然是用动态规划来预处理.这部分依据不同的题目采用不同的写法,通常是一个线性dp.比如此题中,题目要求的是各数位之和在一个区间内,那么我们的dp就可以维护成一个前缀和数组.当然在这题中没有用前缀和优化的这个必要,因为时间复杂度完全够.我们就朴素的定义一个三维的状态:

\(dp[i][j][k]\) 表示当前数一共有i位,最高位为j(这两个信息一般是一定有的,因为第一部分中需要使用这两个状态来计算),当前各数位的和为k的数的个数.

这个很好写出状态转移方程,因为当前求和为k,那就是前一位求和为k - j,只需要枚举前一位是几即可,即:

$ dp[i][j][k] = \sum_{l = 0}^{9}dp[i - 1][l][k - j] $

注意累加的结果可能很大,过程中要进行取模.

这两部分分别写好之后,只需要考虑如何在第一部分中具体的直接使用第二部分的结果即可.显然,只需要把当前的位数和最高位直接带进去,然后对在一个范围内的k所对应的dp数组的值进行求和即可.这里的最高位从0枚举到上界这一位-1,分别加到总情况数中.注意到每次循环的进行其实都是在往二叉树的右支走,所以k的这个范围就是将题目给的范围减去上界当前位之前的各位之和即可.当区间不存在时,意味着当前这个右支及以后的所有都不能到达,直接返回即可.如果完整走完了二叉树的所有右支,说明这也是一种合法情况,需要将总情况数加一(前面计算的均为左支,此处是右支,需要单独加).

另外的一些具体细节可以自行看代码.当然,这份代码的时间复杂度并不是最优的(瓶颈主要在第二部分),此处仅作为参考.采用其他的数位dp写法或对此处的dp采用一些技巧进行优化可能能够获得更优的时间复杂度.

参考代码

class Solution {
public:
    long long dp[25][10][410];
    
    // 预处理dp数组(第二部分)
    Solution() {
        memset(dp,0,sizeof(dp));

        // 基线条件,当前只有一位,单独计算
        for(int i = 0;i <= 9;i++) {
            dp[1][i][i] = 1;
        }

        // 线性dp,直接递推即可
        for(int i = 2;i < 25;i++) {
            for(int j = 0;j <= 9;j++) {
                for(int k = j;k < 410 - j;k++) {
                    for(int l = 0;l <= 9;l++) {
                        dp[i][j][k] =(dp[i][j][k] + dp[i - 1][l][k - j]) % 1000000007;
                    }
                }
            }
        }
    }

    int count(string num1, string num2, int min_sum, int max_sum) {
        // 将最高位放在末尾,方便取dp数组中的值
        reverse(num1.begin(),num1.end());
        reverse(num2.begin(),num2.end());

        // 计算给定上界时的数量
        auto cnt = [&](string num) -> long long {
            // 此法中,单独一个0需要特判
            if(num == "0") return 0;

            // 结果
            long long res = 0;
            // 记录走到当前结点前,各数位的和(因为走的右支,所以前面的各数位就是上界中对应位置上的数)
            long long last = 0;

            // 遍历各位
            for(int i = num.size() - 1;i >= 0;i--) {
                int x = num[i] - '0';
                
                // 左支存在时可以直接计算,当前位分别取0到x-1
                for(int k = 0;k < x;k++) {
                    // 当前区间就是要求的区间减掉已经累加的各位数的和
                    for(int j = max(0ll,min_sum - last);j <= max_sum - last;j++) {
                        res = (res + dp[i + 1][k][j]) % 1000000007;
                    }
                }

                // 左支已经直接计算完毕,接下来的递归均为向右支走,累加当前数位
                last += x;
                if(max_sum - last < 0) break; // 不能再往右支走了,结束递推

                // 右支成功走到底,增加这一种情况
                if(!i && last >= min_sum && last <= max_sum) {
                    res = (res + 1) % 1000000007;
                } 
            }
            return res;
        };

        // 高精度减一
        for(auto &c : num1) {
            if(c == '0') {
                c = '9';
            } else {
                c -= 1;
                break;
            }
        }

        // 计算结果,注意取模防负
        return (cnt(num2) - cnt(num1) + 1000000007) % 1000000007;
    }
};

"正是我们每天反复做的事情，最终造就了我们，优秀不是一种行为，而是一种习惯" ---亚里士多德