https://leetcode.com/problems/numbers-at-most-n-given-digit-set

This is one of those numbers problems, where we need to go digit by digit and count number of good number. Let me explain what I mean by that: imagine, that we have number n = 314792 and digits = [1,3,5,7], then we can split all numbers which are less or equal to n to several groups.

  1. n itself: 314792.
  2. Numbers, which have form 31479?, where question mark can be any number less than 2 and in our digits array.
  3. Numbers, which have form 3147?*, where question mark can be any number less than 9 and * mark means any digit from digits.
  4. Numbers, which have form 314?**, where question mark can be any number less than 7 and * again means any digit from digits.
  5. Numbers, which have form 31?***.
  6. Numbers, which have form 3?****.
  7. Numbers, which have form ?*****.
  8. Finally, numbers which have form *****, ****, ***, **, *.

We can note, that we need to evaluate several times number of digits in digits which are less than some digit 2, 9, 7 and so on. We can just precompute array up for this purpuse. Also let T be number of our digits, str_n will be string created from n, k be length of this string and d_set be set of all digits.

What we need to do in main part of my algorithm:

  1. Iterate from digit to digit, starting with i = 0.
  2. Check if number constructed from first i digits have only digits inside: see step 4: there we have numbers, starting with 314 and they can not be good. On step 5 however numbers start with 31 and it is OK for us.
  3. Add up[int(str_n[i])] multiplied by T to the power k-i-1 to our answer: just combinatorics to choose candidates for ? place and for * place.
  4. Finally, we need to deal with the last case: numbers with less number of digits and with first case: number n itself.

Complexity: time complexity is O(10k + T), where k is lenght of our number n and k is lenght of digits. Actually, it can be just written as O(T), because 10k <= 100 always. Space complexity is O(T).

class Solution:
    def atMostNGivenDigitSet(self, digits, n):
        up, ans, T, str_n = [0]*10, 0, len(digits), str(n)
        for i in range(10):
            for dig in digits:
                up[i] += (dig < str(i))
        
        k, d_set = len(str_n), set(digits)
        for i in range(k):
            if i > 0 and str_n[i-1] not in d_set: break
            ans += up[int(str_n[i])] * T**(k-i-1)
        
        addon = (T**k - 1)//(T-1) - 1 if T != 1 else k - 1
        return ans + addon + (not set(str_n) - set(digits))

If you like the solution, you can upvote it on leetcode discussion section: Problem 0902