Digitwise加法K次后数字位数计算超时问题求助

你的代码超时的核心原因是每次循环都需要处理整个超大数字的字符串/数组转换：随着K增大，数字的位数会快速增长（每次分裂都会增加位数），导致每一轮循环的时间成本越来越高，最终时间复杂度达到O(K*M)（M为每一步的数字位数），远超出题目要求的O(K logn)限制。

优化思路 #

我们不需要维护完整的数字，只需要跟踪每一位数字的分裂次数：

• 当某一位是9时，执行digitwise add会分裂为10，使总位数+1；
• 其他数字加1仅改变数值，不会增加位数。

通过动态规划预计算单个数字（1和0）在m次操作中的分裂次数，再推导初始每一位数字的总分裂次数，最终总位数 = 初始位数 + 总分裂次数，结果取模1_000_000_007即可。

优化后的代码 #

MOD = 10**9 + 7

def digitwise_addition(n, K):
    # 将初始数字拆分为单个数字的列表
    digits = list(map(int, str(n)))
    initial_length = len(digits)
    
    if K == 0:
        return initial_length % MOD
    
    # 预计算f1和f0数组：
    # f1[m] = 数字1经过m次digitwise add后的分裂次数
    # f0[m] = 数字0经过m次digitwise add后的分裂次数
    f1 = [0] * (K + 1)
    f0 = [0] * (K + 1)
    
    for m in range(9, K + 1):
        f1[m] = (1 + f1[m - 9] + f0[m - 9]) % MOD
    
    for m in range(10, K + 1):
        f0[m] = (1 + f1[m - 10] + f0[m - 10]) % MOD
    
    total_splits = 0
    for d in digits:
        if d < 9:
            # 从d变为9需要t次操作，第t+1次操作才会分裂
            t = 9 - d
            if K <= t:
                continue
            remaining = K - (t + 1)
            splits = (1 + f1[remaining] + f0[remaining]) % MOD
            total_splits = (total_splits + splits) % MOD
        else:
            # 初始为9，第一次操作就分裂
            remaining = K - 1
            if remaining < 0:
                continue
            splits = (1 + f1[remaining] + f0[remaining]) % MOD
            total_splits = (total_splits + splits) % MOD
    
    final_length = (initial_length + total_splits) % MOD
    return final_length

代码解释 #

1. 初始处理：将输入数字n拆分为单个数字的列表，记录初始位数。
2. 动态规划预计算：
   • f1[m]：数字1需要8次操作变为9，第9次操作分裂为10，之后每次分裂产生的1和0的分裂次数递归累加，因此递推式为f1[m] = 1 + f1[m-9] + f0[m-9]（m≥9）。
   • f0[m]：数字0需要9次操作变为9，第10次操作分裂为10，递推式为f0[m] = 1 + f1[m-10] + f0[m-10]（m≥10）。
3. 计算总分裂次数：
   • 对于初始数字d<9：先判断K是否足够让d变为9并分裂，若足够则计算剩余操作次数对应的分裂次数。
   • 对于初始数字d=9：第一次操作即分裂，计算剩余操作次数对应的分裂次数。
4. 最终结果：总位数为初始位数加上所有分裂次数之和，取模后返回。

该方案的时间复杂度为O(K + logn)，完全符合题目要求的O(K logn)限制，且避免了处理超大数字的开销，不会出现超时问题。

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