902. Numbers At Most N Given Digit Set

Numbers At Most N Given Digit Set - LeetCodeLeetCode

题目要凑比指定数N小的数字，可以分成两个子问题：1.位数比N少的数字有多少（位数比它少 => 一定比它小）；2.位数和N相同、但比它小的数字有多少；

问题1就是个排列计算，问题2可以dp，状态方程很容易看出来：

f(N) = D中比N[0]小的数字*排列数(n-1/d) + f(N[1:])
where n = len(N), d = len(D)

当N缩到9以下的时候，直接返回D中比N小的digit数量就结束。

需要注意的 edge case 是遇到0。

class Solution:
    def atMostNGivenDigitSet(self, D: List[str], N: int) -> int:
        res = 0
        d, n = len(D), len(str(N))
        for i in range(1, n):
            res += d**(i)
        
        memo = dict()
        
        def countSmaller(number):
            if number in memo.keys(): return memo[number]
            if number <= 9:
                memo[number] = len([x for x in D if int(x) <= number])
                return memo[number]
            head, tail = str(number)[0], str(number)[1:]
            sHead = len([x for x in D if int(x) < int(head)])
            sTail = countSmaller(int(tail)) if tail[0] != '0' and head in D else 0
            memo[number] = sHead*(d**(len(str(number))-1)) + sTail
            return memo[number]
        
        res += countSmaller(N)
        
        return res