Problem statement

https://leetcode.com/problems/restore-the-array/

Solution

The idea is to use dp, where dp[i] is number of ways to split array s[i:]. Then we can start from the end and try to add digit by digit while it is in range [1, k].

Complexity

It is O(n * log(k)) for time and O(n) for space.

Code

class Solution:
    def numberOfArrays(self, s, k):
        n = len(s)
        s = [int(i) for i in s] + [float("inf")]
        dp = [0] * n + [1]
        for i in range(n - 1, -1, -1):
            num, j = s[i], i + 1
            while 1 <= num <= k and j < n + 1:
                dp[i] = (dp[i] + dp[j]) % 1000000007
                num = 10 * num + s[j]
                j += 1
        return dp[0]