小练一下数位 dp

题意

给定一个按 非递减顺序 排列的数字数组 digits 。你可以用任意次数 digits[i] 来写的数字。例如，如果 digits = [‘1’,‘3’,‘5’]，我们可以写数字，如 ‘13’, ‘551’, 和 ‘1351315’。

返回 可以生成的小于或等于给定整数 n 的正整数的个数

样例输入输出

输入：digits = [“1”,”3”,”5”,”7”], n = 100
输出：20
解释：
可写出的 20 个数字是：
1, 3, 5, 7, 11, 13, 15, 17, 31, 33, 35, 37, 51, 53, 55, 57, 71, 73, 75, 77.

分析

以这个样例为例数一下。
对于两位数和一位数，都可以任选, 4 * 4 + 4 = 20
对于 3 位数，第一位只能选 1，, 但是第二位开始没得选了，因此不存在符合条件的三位数。

用 count 函数表示从第 i 位开始枚举有几种做法。
其中有几个参数, nums 表示数位, 高位在 0, 低位在最后, digits 表示可供选择的数位, start 表示当前从第几位开始枚举, lim 表示上一位是否顶到极限。(如果上一位没顶到极限，那么后面的位都可以随便选，否则最多只选到nums[i])

只需要枚举所有第 j 位数即可, 其中 j < nums.size(), 当数位小于 nums.size 时, 符合条件的就是 digits.size() ^ (剩下需要枚举的位数), 这个可以用快速幂来求。

代码

class Solution {
public:
    int atMostNGivenDigitSet(vector<string>& digits, int n) {
        vector<int> nums;
        while(n)nums.push_back(n % 10), n /= 10;
        int ans = 0;
        reverse(nums.begin(), nums.end());
        for(int i = 0; i < nums.size(); i ++){
            ans += count(nums, digits, !i, i);
        }

        return ans;
    }
    int count(vector<int>& nums, vector<string>& digits, bool lim,int start){
        if(!lim){
            return qpow(digits.size(), nums.size()-start);
        }
        if(start >= nums.size())return 1;
        int ans = 0;
        for(int i = 0; i < digits.size(); i ++){
            if(digits[i][0] - '0' == nums[start])
                ans += count(nums, digits, lim, start+1);
            else if(digits[i][0] - '0' < nums[start]) 
                ans += count(nums, digits, false, start+1);
            else break;
        }

        return ans;
    }

    long long qpow(long long a, int b){
        long long res = 1;
        while(b){
            if(b & 1)res = res * a;
            a *= a;
            b >>= 1;
        }

        return res;
    }
};