Skip to content

2698. Find the Punishment Number of an Integer 👍

  • Time: $O(n^2)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
 public:
  int punishmentNumber(int n) {
    int ans = 0;
    for (int i = 1; i <= n; ++i)
      if (isPossible(0, 0, to_string(i * i), 0, i))
        ans += i * i;
    return ans;
  }

 private:
  // Returns true if the sum of any split of `numChars` equals to the target.
  bool isPossible(int accumulate, int running, const string& numChars, int s,
                  int target) {
    if (s == numChars.length())
      return target == accumulate + running;
    const int d = numChars[s] - '0';
    return isPossible(accumulate, running * 10 + d, numChars, s + 1, target) ||
           isPossible(accumulate + running, d, numChars, s + 1, target);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
public class Solution {
  public int punishmentNumber(int n) {
    int ans = 0;
    for (int i = 1; i <= n; ++i)
      if (isPossible(0, 0, String.valueOf(i * i), 0, i))
        ans += i * i;

    return ans;
  }

  // Returns true if the sum of any split of `numChars` equals to the target.
  private boolean isPossible(int accumulate, int running, String numChars, int s, int target) {
    if (s == numChars.length())
      return target == accumulate + running;
    final int d = numChars.charAt(s) - '0';
    return isPossible(accumulate, running * 10 + d, numChars, s + 1, target) ||
        isPossible(accumulate + running, d, numChars, s + 1, target);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Solution:
  def punishmentNumber(self, n: int) -> int:
    def isPossible(accumulate: int, running: int, numChars: List[str], s: int, target: int) -> bool:
      """
      Returns True if the sum of any split of `numChars` equals to the target.
      """
      if s == len(numChars):
        return target == accumulate + running
      d = int(numChars[s])
      return (
          # Keep growing `running`.
          isPossible(accumulate, running * 10 + d, numChars, s + 1, target) or
          # Start a new `running`.
          isPossible(accumulate + running, d, numChars, s + 1, target)
      )

    return sum(i * i
               for i in range(1, n + 1)
               if isPossible(0, 0, str(i * i), 0, i))