【Leetcode】1416. 恢复数组——1920

题目

1416. 恢复数组

某个程序本来应该输出一个整数数组。但是这个程序忘记输出空格了以致输出了一个数字字符串，我们所知道的信息只有：数组中所有整数都在 [1, k] 之间，且数组中的数字都没有前导 0 。

给你字符串 s 和整数 k 。可能会有多种不同的数组恢复结果。

按照上述程序，请你返回所有可能输出字符串 s 的数组方案数。

由于数组方案数可能会很大，请你返回它对 10^9 + 7 取余 后的结果。

• \(1\leq s.length\leq 10^5\)
• \(s\)只包含数字且不包含前导0
• \(1\leq k\leq 10^9\)

思路

一个较为典型的DP，状态转移方程
\(dp[i]=sum(dp[j])\ if\ 1<s[j:i+1]<=k\)
为了方便起见，设置一个哨兵
\(dp[-1]=0\)

如果单单按照上述的状态转移方程求解，那么可以看到时间复杂度是\(O(n^2)\)，但是可以发现的是，由于设置了每次分组的数字上限也就是k，所以在倒序更新dp[i]的时候，如果数字已经超过了k的位数或者比k大了，就没有必要继续遍历了。

而题目给定k的数据范围是\(1\leq k\leq 10^9\)，因此每次倒序更新dp的时候最多不会超过k的长度次循环，即k+1次。

因此时间复杂度为\(O(c\ n）\)

class Solution:
    def numberOfArrays(self, s: str, k: int) -> int:
        MOD = 10 ** 9 + 7
        n = len(s)
        dp = [0] * (n + 1)
        dp[0] = 1
        m = len(str(k))
        for i, x in enumerate(s):
            num = 0
            for j in range(i, -1, -1):
                # 数字位数超过设定数字位数
                if i + 1 - j > m: break
                # 当前位置组成了前导0 或 当前位置之前不能完成分组
                if ord(s[j]) == ord('0') or dp[j] == 0: continue        
                # 更新数字
                num += (ord(s[j]) - ord('0')) * (10 ** (i - j))
                # 超过数字上限即循环停止
                if num > k: break
                dp[i + 1] += dp[j]
            dp[i+1]%=MOD
        return dp[-1]