Numbers At Most N Given Digit Set LeetCode solution in C++

1
class Solution {
2
 public:
3
  using ll = long long;
4

5
  int countLen(int n) {
6
    int cnt = 0;
7
    while (n) {
8
      cnt++;
9
      n /= 10;
10
    }
11
    return cnt;
12
  }
13

14
  vector<int> getDigits(int n) {
15
    vector<int> digits;
16
    while (n) {
17
      digits.push_back(n % 10);
18
      n /= 10;
19
    }
20

21
    reverse(digits.begin(), digits.end());
22
    return digits;
23
  }
24

25
  int solve(int idx, bool flag, vector<int> &digits, vector<int> &number) {
26
    if (!flag) {
27
      int res = 1;
28
      for (int i = idx; i < number.size(); i++) {
29
        res *= digits.size();
30
      }
31
      return res;
32
    }
33

34
    if (idx >= number.size()) return 1;
35

36
    int curr = 0;
37

38
    for (int dig : digits) {
39
      if (dig < number[idx]) {
40
        curr += solve(idx + 1, false, digits, number);
41
      } else if (dig == number[idx]) {
42
        curr += solve(idx + 1, true, digits, number);
43
      } else
44
        break;
45
    }
46

47
    return curr;
48
  }
49

50
  int atMostNGivenDigitSet(vector<string> &digits, int n) {
51
    int len = countLen(n);
52
    vector<int> nums;
53
    for (string str : digits) {
54
      nums.push_back(stoi(str));
55
    }
56
    sort(nums.begin(), nums.end());
57
    ll res = 0;
58

59
    for (int i = 1; i < len; i++) {
60
      int curr = 1;
61
      bool flag = false;
62
      for (int j = 0; j < i; j++) {
63
        flag = true;
64
        curr *= (int)digits.size();
65
      }
66
      if (flag) res += curr;
67
    }
68

69
    vector<int> number = getDigits(n);
70

71
    res += solve(0, true, nums, number);
72

73
    return res;
74
  }
75
};