Count Stepping Numbers in Range

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
MOD = int(1e9) + 7
class Solution:
    def _count(self, n: str) -> int:  # count number of stepping numbers in range [0...n]
        @cache
        def dp(i, tight, lastDigit, leadingZero):
            if i == len(n): 
                return 1  # Found a good number
            maxDigit = int(n[i]) if tight else 9
            ans = 0
            for d in range(maxDigit + 1):
                nxtTight = tight and d == maxDigit
                nxtLeadingZero = leadingZero and d == 0
                if nxtLeadingZero:  # for leading zero, we shouldn't treat lastDigit=d
                    ans = (ans + dp(i + 1, nxtTight, lastDigit, nxtLeadingZero)) % MOD
                elif lastDigit == -1 or abs(lastDigit - d) == 1:
                    ans = (ans + dp(i + 1, nxtTight, d, nxtLeadingZero)) % MOD
            return ans

        return dp(0, True, -1, True)

    def _minusOne(self, s): # s is a string representing a positive integer
        num = int(s) - 1
        return str(num)

    def countSteppingNumbers(self, low: str, high: str) -> int:
        return (self._count(high) - self._count(self._minusOne(low)) + MOD) % MOD