LeetCode - 数位 DP

Posted on Edited on In

See 数位 DP 通用模板，附题单（Python/Java/C++/Go）.

233. Number of Digit One

1
2
3
4
5
6
7
8
9
10
11
12
13
class Solution:
    def countDigitOne(self, n: int) -> int:
        s = str(n)
        @cache
        def f(i: int, cnt1: int, is_limit: bool) -> int:
            if i == len(s):
                return cnt1
            res = 0
            up = int(s[i]) if is_limit else 9
            for d in range(up + 1):
                res += f(i + 1, cnt1 + (d == 1), is_limit and d == up)
            return res
        return f(0, 0, True)

2719. Count of Integers

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Solution:
    def count(self, num1: str, num2: str, min_sum: int, max_sum: int) -> int:
        n = len(num2)
        num1 = num1.zfill(n)

        @cache
        def dfs(i: int, s: int, limit_low: bool, limit_high: bool) -> int:
            if s > max_sum:
                return 0
            if i == n:
                return s >= min_sum
            
            lo = int(num1[i]) if limit_low else 0
            hi = int(num2[i]) if limit_high else 9
            res = 0
            for d in range(lo, hi + 1):
                res += dfs(i + 1, s + d, limit_low and d == lo, limit_high and d == hi)
            return res
        
        return dfs(0, 0, True, True) % 1_000_000_007

788. Rotated Digits

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
DIFFS = (0, 0, 1, -1, -1, 1, 1, -1, 0, 1)
class Solution:
    def rotatedDigits(self, n: int) -> int:
        s = str(n)
        @cache
        def f(i: int, has_diff: bool, is_limit: bool) -> int:
            if i == len(s):
                return has_diff
            res = 0
            up = int(s[i]) if is_limit else 9
            for d in range(0, up + 1):
                if DIFFS[d] != -1:
                    res += f(i + 1, has_diff or DIFFS[d], is_limit and d == up)
            return res
        return f(0, False, True)

902. Numbers At Most N Given Digit Set

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def atMostNGivenDigitSet(self, digits: List[str], n: int) -> int:
        s = str(n)
        @cache
        def f(i: int, is_limit: bool, is_num: bool) -> int:
            if i == len(s):
                return int(is_num)
            res = 0
            if not is_num:
                res = f(i + 1, False, False)
            # no need coz digits[i] != 0
            # low = 0 if is_num else 1
            up = s[i] if is_limit else '9'
            for d in digits:
                if d > up:
                    break
                res += f(i + 1, is_limit and d == up, True)
            return res
        return f(0, True, False)

2376. Count Special Integers

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Solution:
    def countSpecialNumbers(self, n: int) -> int:
        s = str(n)
        @cache
        def f(i: int, mask: int, is_limit: bool, is_num: bool) -> int:
            if i == len(s):
                return int(is_num)
            res = 0
            if not is_num:
                res = f(i + 1, mask, False, False)
            low = 0 if is_num else 1
            up = int(s[i]) if is_limit else 9
            for d in range(low, up + 1):
                if (mask >> d & 1) == 0:
                    res += f(i + 1, mask | (1 << d), is_limit and d == up, True)
            return res
        return f(0, 0, True, False)

600. Non-negative Integers without Consecutive Ones

1
2
3
4
5
6
7
8
9
10
11
12
13
class Solution:
    def findIntegers(self, n: int) -> int:
        s = str(bin(n))[2:]
        @cache
        def f(i: int, pre1: bool, is_limit: bool) -> int:
            if i == len(s):
                return 1
            up = int(s[i]) if is_limit else 1
            res = f(i + 1, False, is_limit and up == 0)
            if not pre1 and up == 1:
                res += f(i + 1, True, is_limit)
            return res
        return f(0, False, True)

1012. Numbers With Repeated Digits

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Solution:
    def numDupDigitsAtMostN(self, n: int) -> int:
        s = str(n)
        @cache
        def f(i: int, mask: int, is_limit: bool, is_num: bool) -> int:
            if i == len(s):
                return int(is_num)
            res = 0
            if not is_num:
                res = f(i + 1, mask, False, False)
            low = 0 if is_num else 1
            up = int(s[i]) if is_limit else 9
            for d in range(low, up + 1):
                if (mask >> d & 1) == 0:
                    res += f(i + 1, mask | (1 << d), is_limit and d == up, True)
            return res
        return n - f(0, 0, True, False)

357. Count Numbers with Unique Digits

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Solution:
    def countNumbersWithUniqueDigits(self, n: int) -> int:
        s = str(pow(10, n))
        @cache
        def f(i: int, mask: int, is_limit: bool, is_num: bool) -> int:
            if i == len(s):
                return int(is_num)
            res = 0
            if not is_num:
                res = f(i + 1, mask, False, False)
            lo = 0 if is_num else 1
            hi = int(s[i]) if is_limit else 9
            for d in range(lo, hi + 1):
                if (mask >> d & 1) == 0:
                    res += f(i + 1, mask | (1 << d), is_limit and d == hi, True)
            return res
        return f(0, 0, True, False) - (1 if len(set(s)) == len(s) else 0) + 1

3007. Maximum Number That Sum of the Prices Is Less Than or Equal to K

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Solution:
    def findMaximumNumber(self, k: int, x: int) -> int:
        def count(num: int) -> int:
            @cache
            def dfs(i: int, cnt1: int, is_limit: bool) -> int:
                if i == 0:
                    return cnt1
                res = 0
                up = num >> (i - 1) & 1 if is_limit else 1
                for d in range(up + 1):
                    res += dfs(i - 1, cnt1 + (d == 1 and i % x == 0), is_limit and d == up)
                return res
            return dfs(num.bit_length(), 0, True)

        return bisect_left(range((k + 1) << x), k + 1, key=count) - 1

2827. Number of Beautiful Integers in the Range

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class Solution:
    def numberOfBeautifulIntegers(self, low: int, high: int, k: int) -> int:
        def calc(high: int) -> int:
            s = str(high)
            @cache
            def dfs(i: int, val: int, diff: int, is_limit: bool, is_num: bool) -> int:
                if i == len(s):
                    return int(is_num and val == 0 and diff == 0)
                res = 0
                if not is_num:
                    res = dfs(i + 1, val, diff, False, False)
                d0 = 0 if is_num else 1
                up = int(s[i]) if is_limit else 9
                for d in range(d0, up + 1):
                    res += dfs(i + 1, (val * 10 + d) % k, diff + d % 2 * 2 - 1, is_limit and d == up, True)
                return res
            return dfs(0, 0, 0, True, False)
        return calc(high) - calc(low - 1)

2999. Count the Number of Powerful Integers

1

2801. Count Stepping Numbers in Range

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def countSteppingNumbers(self, low: str, high: str) -> int:
        MOD = 1_000_000_007
        def calc(s: str) -> int:
            @cache
            def f(i: int, pre: int, is_limit: bool, is_num: bool) -> int:
                if i == len(s):
                    return 1
                res = 0
                if not is_num:
                    res = f(i + 1, -1, False, False)
                lo = 0 if is_num else 1
                hi = int(s[i]) if is_limit else 9
                for d in range(lo, hi + 1):
                    if not is_num or abs(d - pre) == 1:
                        res += f(i + 1, d, is_limit and d == hi, True)
                return res % MOD
            return f(0, -1, True, False)
        return (calc(high) - calc(str(int(low) - 1))) % MOD