Skip to content

2801. Count Stepping Numbers in Range 👍

  • Time: $O(|\texttt{high}| \cdot 10 \cdot 2^2 \cdot 10)$
  • Space: $O(|\texttt{high}| \cdot 10 \cdot 2^2)$
 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
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class Solution {
 public:
  int countSteppingNumbers(string low, string high) {
    const string lowWithLeadingZeros =
        string(high.length() - low.length(), '0') + low;
    vector<vector<vector<vector<int>>>> mem(
        high.length(), vector<vector<vector<int>>>(
                           11, vector<vector<int>>(2, vector<int>(2, -1))));
    return count(lowWithLeadingZeros, high, 0, 10, /*isLeadingZero=*/true, true,
                 true, mem);
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  // Returns the number of valid integers, considering the i-th digit, where
  // `prevDigit` is the previous digit, `isTight1` indicates if the current
  // digit is tightly bound for `low`, and `isTight2` indicates if the current
  // digit is tightly bound for `high`.
  int count(const string& low, const string& high, int i, int prevDigit,
            bool isLeadingZero, bool isTight1, bool isTight2,
            vector<vector<vector<vector<int>>>>& mem) {
    if (i == high.length())
      return 1;
    if (mem[i][prevDigit][isTight1][isTight2] != -1)
      return mem[i][prevDigit][isTight1][isTight2];

    int res = 0;
    const int minDigit = isTight1 ? low[i] - '0' : 0;
    const int maxDigit = isTight2 ? high[i] - '0' : 9;

    for (int d = minDigit; d <= maxDigit; ++d) {
      const bool nextIsTight1 = isTight1 && (d == minDigit);
      const bool nextIsTight2 = isTight2 && (d == maxDigit);
      if (isLeadingZero)
        // Can place any digit in [minDigit, maxDigit].
        res += count(low, high, i + 1, d, isLeadingZero && d == 0, nextIsTight1,
                     nextIsTight2, mem);
      else if (abs(d - prevDigit) == 1)
        // Can only place prevDigit - 1 or prevDigit + 1.
        res +=
            count(low, high, i + 1, d, false, nextIsTight1, nextIsTight2, mem);
      res %= kMod;
    }

    return mem[i][prevDigit][isTight1][isTight2] = res;
  }
};
 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
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class Solution {
  public int countSteppingNumbers(String low, String high) {
    final String lowWithLeadingZeros =
        String.valueOf('0').repeat(high.length() - low.length()) + low;
    Integer[][][][] mem = new Integer[high.length()][11][2][2];
    return count(lowWithLeadingZeros, high, 0, 10, /*isLeadingZero=*/true, true, true, mem);
  }

  private static final int kMod = 1_000_000_007;

  // Returns the number of valid integers, considering the i-th digit, where
  // `prevDigit` is the previous digit, `isTight1` indicates if the current
  // digit is tightly bound for `low`, and `isTight2` indicates if the current
  // digit is tightly bound for `high`.
  private int count(final String low, final String high, int i, int prevDigit,
                    boolean isLeadingZero, boolean isTight1, boolean isTight2,
                    Integer[][][][] mem) {
    if (i == high.length())
      return 1;
    if (mem[i][prevDigit][isTight1 ? 1 : 0][isTight2 ? 1 : 0] != null)
      return mem[i][prevDigit][isTight1 ? 1 : 0][isTight2 ? 1 : 0];

    int res = 0;
    final int minDigit = isTight1 ? low.charAt(i) - '0' : 0;
    final int maxDigit = isTight2 ? high.charAt(i) - '0' : 9;

    for (int d = minDigit; d <= maxDigit; ++d) {
      final boolean nextIsTight1 = isTight1 && (d == minDigit);
      final boolean nextIsTight2 = isTight2 && (d == maxDigit);
      if (isLeadingZero)
        // Can place any digit in [minDigit, maxDigit].
        res += count(low, high, i + 1, d, isLeadingZero && d == 0, nextIsTight1, nextIsTight2, mem);
      else if (Math.abs(d - prevDigit) == 1)
        // Can only place prevDigit - 1 or prevDigit + 1.
        res += count(low, high, i + 1, d, false, nextIsTight1, nextIsTight2, mem);
      res %= kMod;
    }

    return mem[i][prevDigit][isTight1 ? 1 : 0][isTight2 ? 1 : 0] = res;
  }
}
 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
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class Solution:
  def countSteppingNumbers(self, low: str, high: str) -> int:
    kMod = 1_000_000_007
    low = '0' * (len(high) - len(low)) + low

    @functools.lru_cache(None)
    def dp(
        i: int,
        prevDigit: int,
        isLeadingZero: bool,
        isTight1: bool,
        isTight2: bool,
    ) -> int:
      """
      Returns the number of valid integers, considering the i-th digit, where
      `prevDigit` is the previous digit, `isTight1` indicates if the current
      digit is tightly bound for `low`, and `isTight2` indicates if the current
      digit is tightly bound for `high`.
      """
      if i == len(high):
        return 1

      res = 0
      minDigit = int(low[i]) if isTight1 else 0
      maxDigit = int(high[i]) if isTight2 else 9

      for d in range(minDigit, maxDigit + 1):
        nextIsTight1 = isTight1 and (d == minDigit)
        nextIsTight2 = isTight2 and (d == maxDigit)
        if isLeadingZero:
          # Can place any digit in [minDigit, maxDigit].
          res += dp(i + 1, d, isLeadingZero and d ==
                    0, nextIsTight1, nextIsTight2)
        elif abs(d - prevDigit) == 1:
          res += dp(i + 1, d, False, nextIsTight1, nextIsTight2)
        res %= kMod

      return res

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