Problem Informations

• Problem Index: 2698
• Problem Link: Find the Punishment Number of an Integer
• Topics: Math, Backtracking

Problem Description

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 \leq i \leq n$
• The decimal representation of $i \times 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 \leq n \leq 1000$

Intuition

The problem involves determining specific integers within a given range that satisfy a particular partition condition when squared. Specifically, for a given integer $i$, its square $i \times i$ should be splittable into contiguous substrings such that the sum of these substrings equals $i$ itself. The goal is to sum the squares of all such integers from 1 to $n$. This boils down to a combination of digit manipulation and dynamic programming to efficiently explore possible partitions.

Approach

The solution leverages digit extraction and a dynamic programming approach to determine if a given integer $i$ meets the specified condition. Below is the step-by-step approach:

1. Digit Counting: A helper function calculates the number of digits in a number, aiding in iterating through possible substring partitions.

2. Substring Extraction: Another helper function allows us to extract substring numbers from any given number using digit indices. This is critical for partitioning the squared number effectively.

3. Partition Validation: A dynamic programming approach is employed to explore all possible ways to partition the decimal representation of $i \times i$. The process is as follows:

   • Initialize a DP table dp where each entry dp[k] contains all possible sums up to the $k$-th digit position of the number being squared.
   • Populate this table by iterating through potential end indices for the substring and calculate all possible partition sums extending from earlier partitions up to this index.
   • Check if any sum in the final list of partitions equals the original number $i$. If found, this indicates a valid partitioning.

4. Main Execution: The main function punishmentNumber iterates through all integers from 1 to $n`:

   • For each integer $i$, calculate its square.
   • Use the partition validation function to check if this square can be partitioned to sum to $i`.
   • Accumulate the square to the total punishment number if the partition condition is satisfied.

This approach efficiently checks each possible number for its square’s partition property, ensuring that each valid case contributes its square to the final sum. The combination of precise digit operations and dynamic programming makes the solution both comprehensive and efficient within the given constraints.

Code

C++

class Solution {
public:
    // Function to return the number of digits in a given number
    int countDigits(int num) {
        int count = 0;
        while (num > 0) {
            num /= 10;
            count++;
        }
        return count;
    }

    // Function to extract a substring number from 'num'
    // starting at 'beg_idx' to 'end_idx' given that 'num' has 'digits' digits
    int extractSubNumber(int num, int beg_idx, int end_idx, int digits) {
        // Remove leading digits that are beyond end_idx
        num /= pow(10, digits - end_idx - 1);
        // Extract the relevant substring part as a number
        return num % int(pow(10, end_idx - beg_idx + 1));
    }

    // Function that checks if the square of a number can satisfy the condition
    bool isValidPartition(int num, int target) {
        int digits = countDigits(num);
        
        // Dynamic programming table to store possible sum combinations up to each index
        vector<vector<int>> dp(digits);
        
        // Initialize the first digit
        dp[0].push_back(extractSubNumber(num, 0, 0, digits));
        
        // Fill the DP table with all possible sums for each substring end index
        for (int i = 1; i < digits; i++) {
            dp[i].push_back(extractSubNumber(num, 0, i, digits));
            for (int j = 1; j <= i; j++) {
                int currentNum = extractSubNumber(num, j, i, digits);
                if (!dp[j - 1].empty()) {
                    for (int partialSum : dp[j - 1]) {
                        dp[i].push_back(partialSum + currentNum);
                    }
                }
            }
        }

        // Check if any of the sums matches the target value
        for (int partitionSum : dp[digits - 1]) {
            if (partitionSum == target) {
                return true;
            }
        }
        
        return false;
    }

    // Main function to calculate the punishment number for 'n'
    int punishmentNumber(int n) {
        int totalPunishment = 0;
        
        // Iterate through all integers from 1 to n
        for (int i = 1; i <= n; i++) {
            int square = i * i;
            // Check if the square of the current number can be partitioned to sum to the number
            if (isValidPartition(square, i)) {
                totalPunishment += square;
            }
        }
        
        return totalPunishment;
    }
};

Python

```python
class Solution:
    # Function to return the number of digits in a given number
    def countDigits(self, num):
        count = 0
        while num > 0:
            num //= 10
            count += 1
        return count

    # Function to extract a substring number from 'num'
    # starting at 'beg_idx' to 'end_idx' given that 'num' has 'digits' digits
    def extractSubNumber(self, num, beg_idx, end_idx, digits):
        # Remove leading digits that are beyond end_idx
        num //= 10**(digits - end_idx - 1)
        # Extract the relevant substring part as a number
        return num % 10**(end_idx - beg_idx + 1)

    # Function that checks if the square of a number can satisfy the condition
    def isValidPartition(self, num, target):
        digits = self.countDigits(num)
        
        # Dynamic programming table to store possible sum combinations up to each index
        dp = [[] for _ in range(digits)]
        
        # Initialize the first digit
        dp[0].append(self.extractSubNumber(num, 0, 0, digits))
        
        # Fill the DP table with all possible sums for each substring end index
        for i in range(1, digits):
            dp[i].append(self.extractSubNumber(num, 0, i, digits))
            for j in range(1, i + 1):
                currentNum = self.extractSubNumber(num, j, i, digits)
                if dp[j - 1]:
                    for partialSum in dp[j - 1]:
                        dp[i].append(partialSum + currentNum)

        # Check if any of the sums matches the target value
        for partitionSum in dp[digits - 1]:
            if partitionSum == target:
                return True
        
        return False

    # Main function to calculate the punishment number for 'n'
    def punishmentNumber(self, n):
        totalPunishment = 0
        
        # Iterate through all integers from 1 to n
        for i in range(1, n + 1):
            square = i * i
            # Check if the square of the current number can be partitioned to sum to the number
            if self.isValidPartition(square, i):
                totalPunishment += square
        
        return totalPunishment

# Complexity
- Time complexity: The time complexity of the `punishmentNumber` method is $O(n \cdot d^3)$, where $n$ is the input integer and $d$ is the number of digits in the square of the maximum number ($n \times n$). This is derived as follows:

  - The outer loop iterates $n$ times, once for each integer from $1$ to $n$.
  - Inside this loop, the `isValidPartition` method is called, which:
    - Determines the number of digits in $i \times i$, which is $O(\log(i^2)) = O(\log(i))$ [accounts for $\log(n^2) \approx 2 \log(n)$ which is capped to $d$]. 
    - Initializes a dynamic programming table with $d$ rows, each potentially requiring evaluation of all possible subpartitions, i.e., from $0$ to $d-1$. This makes the loop complexity roughly $O(d^3)$ considering nested loops and potential DP propagation.
  - Thus, the combination of $n$ outer iterations and $O(d^3)$ inner complexity gives $O(n \cdot d^3)$.

- Space complexity: The space complexity is $O(d^2)$, where $d$ is the number of digits in the largest possible square number ($n \times n$). This is due to the dynamic programming table used in `isValidPartition`, which can store all possible sums up to each digit position in the number, potentially leading to $O(d)$ distinct entries for each digit from $0$ to $d-1$.