328不大于n的好数计数问题

问题描述

小 C 定义一个正整数为“好数”，当且仅当这个数的相邻数位都不相同。比如，1024 和 121 是好数，但 233 不是。现在小 C 想知道，不大于给定整数 n 的好数有多少个。结果对 $10^9 + 7$ 取模。

解题思路

1. 转换为字符串：首先将整数 n 转换为字符串 s，以便逐位处理。

2. 定义状态：

   • i：当前处理到第 i 位。
   • pre：前一位的数字，用于判断当前位是否与前一位不同。
   • is_limit：当前位是否受到上界的限制，即前面的数字是否已经与 n 的对应位相同。
   • is_num：当前是否已经开始构造数字（避免前导零）。

3. 递归函数 dfs：

   • 当 i 等于 len(s) 时，表示所有位都已经处理完毕，如果已经开始构造数字，则返回 1，否则返回 0。
   • 如果尚未开始构造数字，可以选择跳过当前位（即不选择任何数字）。
   • 对于每一位，可以选择从 0 到 9（根据 is_limit 确定上界）。
   • 选择的数字 d 需要与前一位数字 pre 不同。
   • 递归处理下一位，更新状态参数。

4. 记忆化搜索：使用 @cache 装饰器对递归函数进行记忆化，以避免重复计算，提高效率。

5. 结果：最终返回符合条件的“好数”数量，对 10^9 + 7 取模。

代码解析

from functools import cache

def solution(n: int) -> int:
    s = str(n)

    @cache
    def dfs(i, pre, is_limit, is_num):
        if i == len(s):
            return int(is_num)
        res = 0
        if not is_num:
            res = dfs(i + 1, pre, False, False)
        low = 0 if is_num else 1
        up = int(s[i]) if is_limit else 9
        for d in range(low, up + 1):
            if d != pre:
                res += dfs(i + 1, d, is_limit and d == up, True)
        return res

    return dfs(0, -1, True, False)

复杂度分析

• 时间复杂度：O(len(s) * 10 * 10)，其中 len(s) 是整数 n 的位数，每一位有 10 种选择。动态规划的时间复杂度 = 状态个数 × 单个状态的计算时间
• 空间复杂度：O(len(s) * 10)，用于记忆化存储中间结果。