数位DP相关题目及通用模版

python 复制代码
# 数位DP 通用模版
class CountSpecialNumbers:
    """
    根据本题提炼出 数位DP 通用模版
    2376. 统计特殊整数
    https://leetcode.cn/problems/count-special-integers/description/
    """
    def solution(self, n: int) -> int:
        s = str(n)

		# cache 表示记忆化搜索，替代常见的 memo 数组
        @cache
        def f(i: int, mask: int, is_limit: bool, is_num: bool) -> int:
            """
            f(i,isLimit,isNum) 表示构造从左往右第 i 位及其之后数位的合法方案数
            :param i: 当前来到的数位
            :param mask: 前面选过的数字集合
            :param is_limit: 前面的数字是否全部顶格获取
            :param is_num: i 前面的数位是否填了数字
            :return:
            """
            if i == len(s):
                # is_num 为True，表示得到了1个合法数字
                return int(is_num)

            res = 0

            # 说明前面数位还未填充任何数字，mask自然为0
            if not is_num:
                # 可以跳过当前数位
                res = f(i+1, 0, False, False)

            low = 0 if is_num else 1  # 如果前面没有填数字，必须从 1 开始（因为不能有前导零）
            up = int(s[i]) if is_limit else 9  # 如果前面填的数字都和 n 的一样，那么这一位至多填 s[i]（否则就超过 n 啦）
            for d in range(low, up+1):  # 枚举d
                if (mask >> d & 1) == 0:  # d不在mask中
                    res += f(i+1, mask | (1 << d), is_limit and d == up, True)

            return res

        return f(0, 0, True, False)


class AtMostNGivenDigitSet:
    """
    902. 最大为 N 的数字组合
    https://leetcode.cn/problems/numbers-at-most-n-given-digit-set/description/
    """
    def solution(self, digits: List[str], n: int) -> int:
        """

        :param digits:
        :param n:
        :return:
        """
        s = str(n)

        @cache
        def f(i: int, is_limit: bool, is_num: bool) -> int:
            """
            f(i,isLimit,isNum) 表示构造从左往右第 i 位及其之后数位的合法方案数
            :param i: 第i位要填充的数字
            :param is_limit: 当前数字是否都是顶格获取的，比如 7653849 这个数字填充第三位的时候，
            所谓顶格选其实就是前两位选了 76，同理 765， 7653， 7653849 等各个数位都是顶格选的，
            简单理解也可以认为是最大值。
            :param is_num: 表示第i-1位是否填充了数字
            :return:
            """
            if i == len(s):
                return int(is_num)  # 如果填了数字，则为 1 种合法方案

            res = 0
            if not is_num:  # 前面一位没填数字，那么可以跳过当前数位，也不填数字
                # is_limit 改为 False，因为没有填数字，位数都比 n 要短，自然不会受到 n 的约束
                # is_num 仍然为 False，因为没有填任何数字
                res = f(i+1, False, False)

            up = s[i] if is_limit else '9'  # 根据是否受到约束，决定当前位可以填的数字上限

            # 注意：对于一般的题目而言，如果此时 is_num 为 False，则必须从 1 开始枚举（不能第一位是0），
            # 由于本题 digits 没有 0，所以无需处理这种情况
            for d in digits:
                # d 超过上限，由于 digits 是有序的，后面的 d 都会超过上限，故退出循环
                if d > up:
                    break

                # is_limit：如果当前受到 n 的约束，且填的数字等于上限，那么后面仍然会受到 n 的约束
                # is_num 为 True，因为填了数字
                res += f(i+1, is_limit and d == up, True)

            return res

        return f(0, True, False)

    def solution1(self, digits: List[str], n: int) -> int:
        """
        动态规划
        :param digits:
        :param n:
        :return:
        """
        m = len(digits)
        s = str(n)
        k = len(s)

        # 定义dp[i][0] 表示由digits构成且 小于 n的前i位 的数字的个数
        # 定义dp[i][1] 表示由digits构成且 等于 n的前i位 的数字的个数
        dp = [[0, 0] for _ in range(k + 1)]

        dp[0][1] = 1
        for i in range(1, k + 1):
            for d in digits:
                if d == s[i - 1]:
                    dp[i][1] = dp[i - 1][1]
                elif d < s[i - 1]:
                    dp[i][0] += dp[i - 1][1]
                else:
                    break

            if i > 1:
                dp[i][0] += m + dp[i - 1][0] * m

        return sum(dp[k])