# Find the Punishment Number of an Integer
**Difficulty:** MEDIUM
[External](https://leetcode.com/problems/find-the-punishment-number-of-an-integer)
Canonical: https://scaleengineer.com/dsa/problems/find-the-punishment-number-of-an-integer
**Patterns:** [Math](https://scaleengineer.com/dsa/patterns/math), [Backtracking](https://scaleengineer.com/dsa/patterns/backtracking)
---
## Problem
Given a positive integer `n`, return _the **punishment number**_ of `n`.

The **punishment number** of `n` is defined as the sum of the squares of all integers `i` such that:

* `1 <= i <= n`
* The decimal representation of `i * i` can be partitioned into contiguous substrings such that the sum of the integer values of these substrings equals `i`.

**Example 1:**

**Input:** n = 10
**Output:** 182
**Explanation:** There are exactly 3 integers i in the range [1, 10] that satisfy the conditions in the statement:
- 1 since 1 * 1 = 1
- 9 since 9 * 9 = 81 and 81 can be partitioned into 8 and 1 with a sum equal to 8 + 1 == 9.
- 10 since 10 * 10 = 100 and 100 can be partitioned into 10 and 0 with a sum equal to 10 + 0 == 10.
Hence, the punishment number of 10 is 1 + 81 + 100 = 182

**Example 2:**

**Input:** n = 37
**Output:** 1478
**Explanation:** There are exactly 4 integers i in the range [1, 37] that satisfy the conditions in the statement:
- 1 since 1 * 1 = 1. 
- 9 since 9 * 9 = 81 and 81 can be partitioned into 8 + 1. 
- 10 since 10 * 10 = 100 and 100 can be partitioned into 10 + 0. 
- 36 since 36 * 36 = 1296 and 1296 can be partitioned into 1 + 29 + 6.
Hence, the punishment number of 37 is 1 + 81 + 100 + 1296 = 1478

**Constraints:**

* `1 <= n <= 1000`

# Approaches
## Brute-Force with Recursive Backtracking
This approach iterates through each integer `i` from 1 to `n`. For each `i`, it first calculates its square, `i*i`, and converts it into a string. Then, a recursive backtracking function is used to determine if this string can be partitioned into substrings whose integer values sum up to `i`. If a valid partition exists, `i*i` is added to a running total. This method directly translates the problem statement into code.
**Time:** O(N * K), where `N` is the input `n`, and `K` is the cost of the `canPartition` check. The `dfs` function's complexity is exponential in the length of the string `s`, let's say `L`. `L = O(log(i^2)) = O(log i)`. The number of partitions of a string of length `L` is `2^(L-1)`. So, `K` is roughly `O(2^L)`. Given `N <= 1000`, `L` is at most 7, so this is feasible. · **Space:** O(log N). The space complexity is determined by the maximum depth of the recursion stack for the `dfs` function. The depth is at most the length of the string representation of `i*i`, which is proportional to `log(i^2)` or `log N`.
**Pros:** Conceptually simple and a direct implementation of the problem requirements.; Does not require complex data structures.
**Cons:** Can be slow due to the exponential nature of the recursive partition check, especially for larger values of `n` (though feasible for the given constraints).; For multiple test cases (as in a competitive programming platform), it recomputes the same results for smaller `n` values repeatedly.
### Explanation
The main function `punishmentNumber(n)` initializes a sum to zero and then loops from `i = 1` to `n`. Inside the loop, it calculates `square = i * i` and converts it to a `String s`. It then calls a helper function, `canPartition(s, i)`, to check if `s` can be partitioned to sum to `i`. If `canPartition` returns `true`, `square` is added to the sum.

The `canPartition` function itself is a wrapper that initiates a recursive search. It calls a backtracking function `dfs(s, target, index, currentSum)` which explores all possible partitions of the string `s` starting from a given `index` to see if they can sum up to the `target`.

```java
class Solution {
    public int punishmentNumber(int n) {
        int totalPunishment = 0;
        for (int i = 1; i <= n; i++) {
            String s = String.valueOf(i * i);
            if (canPartition(s, i)) {
                totalPunishment += i * i;
            }
        }
        return totalPunishment;
    }

    private boolean canPartition(String s, int target) {
        return dfs(s, target, 0, 0);
    }

    private boolean dfs(String s, int target, int index, int currentSum) {
        if (index == s.length()) {
            return currentSum == target;
        }

        for (int j = index; j < s.length(); j++) {
            String sub = s.substring(index, j + 1);
            int num = Integer.parseInt(sub);

            if (currentSum + num > target) {
                // Optimization: if current sum already exceeds target,
                // no need to check longer numbers from this index.
                break;
            }

            if (dfs(s, target, j + 1, currentSum + num)) {
                return true;
            }
        }

        return false;
    }
}
```
### Algorithm
- Initialize `punishmentSum = 0`.
- For `i` from 1 to `n`:
  - Let `s = String.valueOf(i * i)`.
  - If a helper function `canPartition(s, i)` returns `true`, add `i * i` to `punishmentSum`.
- Return `punishmentSum`.
- The `canPartition` function uses a recursive depth-first search (`dfs`) to check the condition.
- **`dfs(s, target, index, currentSum)` function**:
  - **Base Case**: If `index` reaches the end of the string `s`, it means we have processed the entire string. Return `true` if `currentSum` equals `target`, otherwise `false`.
  - **Recursive Step**: Iterate through all possible next substrings starting from `index`. For each substring:
    - Parse it to an integer `num`.
    - If `currentSum + num` exceeds `target`, we can prune this path and all subsequent paths from this `index` because numbers will only get larger.
    - Make a recursive call `dfs(s, target, j + 1, currentSum + num)` for the rest of the string.
    - If any recursive call returns `true`, it means a valid partition is found, so propagate `true` up the call stack.
  - If the loop completes without finding a valid partition, return `false`.

## Pre-computation with Caching
This approach leverages the constraint that `n` is at most 1000. Instead of calculating the result on-the-fly for each call, we can pre-compute the punishment numbers for all integers from 1 to 1000 and store them in a cache (e.g., an array). The main function then simply looks up the result from this cache. The expensive computation is performed only once, making subsequent calls for any `n <= 1000` instantaneous.
**Time:** O(1) for each call to `punishmentNumber(n)`. The one-time pre-computation cost is a fixed constant, calculated as `sum_{i=1 to 1000} O(2^(log i))`. This upfront cost is paid only once when the class is loaded. · **Space:** O(1) with respect to the input `n`, as the space used is constant. It's `O(MAX_N)` for the cache array (`1001 * 4` bytes) and `O(log(MAX_N))` for the recursion stack during the one-time pre-computation.
**Pros:** Extremely fast query time (`O(1)`) for any given `n` after the initial setup.; Amortizes the computation cost over multiple test cases, making it very efficient in a typical online judge environment.
**Cons:** Requires extra space for the cache, which is `O(MAX_N)`.; Incurs an upfront computation cost, which might be slightly inefficient if the function is guaranteed to be called only once with a very small `n`.
### Explanation
We use a static array, `punishmentCache`, of size 1001 to store the cumulative punishment numbers, where `punishmentCache[i]` will hold the punishment number for `i`. A static block is used to populate this cache when the class is loaded. This block iterates from `i = 1` to 1000, and for each `i`, it applies the same recursive check as the brute-force method. The cumulative sum is maintained in the cache: `punishmentCache[i] = punishmentCache[i-1] + (is_punishable ? i*i : 0)`. The main `punishmentNumber(n)` function, after the one-time initialization, just returns `punishmentCache[n]`. This approach is highly efficient for platforms that run multiple test cases, as the pre-computation cost is amortized.

```java
class Solution {
    private static final int MAX_N = 1000;
    private static final int[] punishmentCache = new int[MAX_N + 1];

    static {
        for (int i = 1; i <= MAX_N; i++) {
            String s = String.valueOf(i * i);
            punishmentCache[i] = punishmentCache[i - 1];
            if (canPartition(s, i)) {
                punishmentCache[i] += i * i;
            }
        }
    }

    private static boolean canPartition(String s, int target) {
        return dfs(s, target, 0, 0);
    }

    private static boolean dfs(String s, int target, int index, int currentSum) {
        if (index == s.length()) {
            return currentSum == target;
        }

        for (int j = index; j < s.length(); j++) {
            String sub = s.substring(index, j + 1);
            int num = Integer.parseInt(sub);
            
            if (currentSum + num > target) {
                break;
            }

            if (dfs(s, target, j + 1, currentSum + num)) {
                return true;
            }
        }
        return false;
    }

    public int punishmentNumber(int n) {
        return punishmentCache[n];
    }
}
```
### Algorithm
- Declare a static integer array `cache` of size 1001 (based on the constraint `n <= 1000`).
- Use a static initializer block to perform a one-time pre-computation.
- Inside the static block:
  - Loop `i` from 1 to 1000.
  - Let `s = String.valueOf(i * i)`.
  - Check if `i` is a punishable number by calling the same `dfs(s, i, 0, 0)` function from the previous approach.
  - Calculate the cumulative punishment number: `cache[i] = cache[i-1]`. If `i` is punishable, add `i*i` to `cache[i]`.
- The main `punishmentNumber(n)` function simply returns the pre-computed value from `cache[n]`.

# Solutions
### Java

```java
class Solution {
public
  int punishmentNumber(int n) {
    int ans = 0;
    for (int i = 1; i <= n; ++i) {
      int x = i * i;
      if (check(x + "", 0, i)) {
        ans += x;
      }
    }
    return ans;
  }
private
  boolean check(String s, int i, int x) {
    int m = s.length();
    if (i >= m) {
      return x == 0;
    }
    int y = 0;
    for (int j = i; j < m; ++j) {
      y = y * 10 + (s.charAt(j) - '0');
      if (y > x) {
        break;
      }
      if (check(s, j + 1, x - y)) {
        return true;
      }
    }
    return false;
  }
}

```

### CPP

```cpp
class Solution {
public:
  int punishmentNumber(int n) {
    int ans = 0;
    for (int i = 1; i <= n; ++i) {
      int x = i * i;
      string s = to_string(x);
      if (check(s, 0, i)) {
        ans += x;
      }
    }
    return ans;
  }
  bool check(const string &s, int i, int x) {
    int m = s.size();
    if (i >= m) {
      return x == 0;
    }
    int y = 0;
    for (int j = i; j < m; ++j) {
      y = y * 10 + s[j] - '0';
      if (y > x) {
        break;
      }
      if (check(s, j + 1, x - y)) {
        return true;
      }
    }
    return false;
  }
};

```

### Python

```python
class Solution:
    def punishmentNumber(self, n: int) -> int: def check(s: str, i: int, x: int) -> bool: m = len(s) if i >= m: return x == 0 y = 0 for j in range(i, m): y = y * 10 + int(s[j]) if y > x: break if check(s, j + 1, x - y): return True return False ans = 0 for i in range(1, n + 1): x = i * i if check(str(x), 0, i): ans += x return ans

```