优化仅用repdigit分解正整数的计数算法（Python）

问题背景 #

计算用repdigit（由单一数字组成的正整数，如1、11、22等）分解正整数n的不同方式数，顺序不同的分解视为不同（例如n=3时有4种分解方式：1+1+1、1+2、2+1、3）。原动态规划解法在n=90000时因超时无法满足要求，需优化。

原代码瓶颈分析 #

原代码采用dp[i] = sum(dp[i - r] for r in rep)的递推逻辑，内层循环遍历所有repdigit。虽然repdigit总数仅约O(log n)个，但Python中循环开销累积，当n=9e4时，总循环次数达4e6级别，导致超时。

优化方案 #

利用Repdigit的数学结构简化递推 #

每个数字d（1-9）对应的repdigit序列为d, 11d, 111d, ...，即d*(10^m -1)/9（m≥1）。我们可以推导该序列对dp[i]的贡献总和s_d[i]的递推公式，避免逐个遍历每个repdigit：

• 当i < d时，s_d[i] = 0（无可用的repdigit）
• 当i ≥ d时，s_d[i] = dp[i - d] + (s_d[i - 10*d] if i ≥ 10*d else 0)

该公式的核心逻辑：最小的repdigit是d，更长的repdigit（如11d、111d）可看作在i-10*d的基础上追加d+10d，因此复用s_d[i-10*d]的结果即可快速计算总和。

重构DP计算逻辑 #

将原内层遍历所有repdigit的操作，替换为遍历9个数字d，通过上述递推式计算s_d[i]并累加得到dp[i]，时间复杂度从O(n*log n)降至O(n)，大幅减少循环次数。

移除冗余操作 #

原代码中的rep.sort()完全多余，递推求和的顺序不影响结果，直接移除即可。

优化后的代码 #

def count_repdigit_decompositions(n: int) -> int:
    dp = [0] * (n + 1)
    dp[0] = 1  # 边界条件：和为0的分解方式有1种（空分解）

    # 为每个数字d维护贡献总和数组s[d]
    s = [[0] * (n + 1) for _ in range(10)]  # s[0]未使用

    for i in range(1, n + 1):
        total = 0
        for d in range(1, 10):
            if i >= d:
                s[d][i] = dp[i - d]
                if i >= 10 * d:
                    s[d][i] += s[d][i - 10 * d]
            total += s[d][i]
        dp[i] = total
    return dp[n]

进一步空间优化（可选） #

若需节省空间，可针对每个d使用滚动变量代替完整数组，仅保存计算所需的历史值。例如对于d=1，只需保存最近10个位置的s值即可，但对于n=9e4来说，原空间占用可忽略，此优化非必需。

内容的提问来源于stack exchange，提问作者Lesserrafim