2698. Find the Punishment Number of an Integer

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

Approach 01:

#include <string>
using namespace std;

class Solution {
public:
    bool canPartition(const string& squareStr, int target, int index, int currentSum) {
        if (index == squareStr.length())
            return currentSum == target;

        int tempSum = 0;
        for (int j = index; j < squareStr.length(); j++) {
            tempSum = tempSum * 10 + (squareStr[j] - '0');

            if (tempSum + currentSum > target)
                break;

            if (canPartition(squareStr, target, j + 1, currentSum + tempSum))
                return true;
        }

        return false;
    }

    int punishmentNumber(int n) {
        int totalSum = 0;

        for (int i = 1; i <= n; i++) {
            string squareStr = to_string(i * i);
            if (canPartition(squareStr, i, 0, 0))
                totalSum += i * i;
        }

        return totalSum;
    }
};

class Solution:
    def punishmentNumber(self, n: int) -> int:
        def can_partition(s, target, index, curr_sum):
            if index == len(s):
                return curr_sum == target
            temp_sum = 0
            for j in range(index, len(s)):
                temp_sum = temp_sum * 10 + int(s[j])
                if temp_sum + curr_sum > target:
                    break
                if can_partition(s, target, j + 1, curr_sum + temp_sum):
                    return True
            return False

        result = 0
        for i in range(1, n + 1):
            square_str = str(i * i)
            if can_partition(square_str, i, 0, 0):
                result += i * i
                
        return result

Time Complexity:

• punishmentNumber(n):

  Iterates from \( 1 \) to \( n \), performing \( O(n) \) operations.

• canPartition(squareStr, target, index, currentSum):

  Uses a recursive approach to check possible partitions of the squared number, leading to an exponential complexity of \( O(2^m) \), where \( m \) is the length of squareStr.

• Overall Time Complexity:

  Since each number up to \( n \) is processed, the worst-case time complexity is \( O(n \cdot 2^m) \).

Space Complexity:

• Recursive Call Stack:

  The depth of the recursion can be at most \( m \), leading to \( O(m) \) auxiliary space.

• String Storage:

  Stores the squared value of each number as a string, contributing \( O(m) \) space per number.

• Overall Space Complexity:

  \( O(m) \) in the worst case.