2698. Find the Punishment Number of an Integer

MediumMathBacktracking
Leetcode Link

Problem Description

The task here is to find the so-called "punishment number" for a given positive integer n. A "punishment number" is essentially the sum of the squares of particular integers ranging from 1 to n. An integer i (within this range) contributes to this sum if and only if the square of i (i * i) can be split into one or more contiguous substrings, and when these substrings are converted to integers and added together, the sum equals i.

For example, if n is 81, we look at each number from 1 to 81, square them, and see if the square can be partitioned into substrings that add up to the original number. In this case, if n is 81, 81 itself does not contribute because 8181 is 6561 which cannot be split into substrings that add to 81. But the number 9 does contribute because 99 is 81, and 8+1 is 9.

Intuition

The solution involves iterating through each number i from 1 to n, calculating i * i, and checking if this square can meet the condition described above.

The intuition behind the implementation is that we can use recursion to check all possible substrings. For every position in the string representation of i * i, we recursively check all ways we could split the remaining string. If we reach the end of the string and the sum of the substrings equals i, then i * i is valid and should be included in our "punishment number" sum.

The check function in the code takes a string s, which represents the square of a number, an index i which indicates where to start partitioning, and an integer x that we are trying to match using the sum of substrings. It tries all possible ways to partition the string from index i onwards and returns true if any of those partitions sum up to x.

The main solution function punishmentNumber goes from 1 to n, checking each number i by calling the check function. If check returns true for a given i, it means we can sum up the substrings of i * i to get i, hence, we add i * i to our cumulative ans which is our punishment number.

This approach efficiently checks all possible partition points in the sequence and uses memorization by the nature of recursion to avoid rechecking the same combination of splits, leading to an optimal solution.

Learn more about Math and Backtracking patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece：

How does quick sort divide the problem into subproblems?

Solution Approach

To solve the "punishment number" problem, the implemented solution leverages recursion and backtracking as its core algorithmic concepts.

The check Helper Function

  • This is a recursive function designed to check if a given squared number can be split into substrings that sum up to the original number before squaring.
  • It takes the string s which represents the square of i, an index i which represents the current position in s that we're examining, and an integer x which is the value of the original number that we're aiming to reconstruct from the substrings.
  • It iterates through the string starting from the current index i and accumulates a value y. If at any point the accumulated y is greater than x, there's no need to continue along that path.
  • If y equals x and we are at the end of the string, then we know this is a valid partitioning. If not, we call check recursively with a new starting position just past the current substring we're considering. This recursive call will return True if the subsequent partitioning works out, or keep trying until all possibilities have been exhausted.

The punishmentNumber Main Function

  • It iterates from 1 to n. For each iteration:
    • Calculates i * i to get the square of the current number.
    • It turns the square into a string and calls the check function with that string, the number 0 for starting index, and the original number i we are trying to partition into.
  • If the check function returns True, which indicates that the square of i can be split into substrings adding up to i, then i * i gets added to the cumulative sum ans.
  • After iteration completes, ans holds the sum of all valid squares and is returned as the final answer.

Algorithm and Data Structures

  • The main data structure used here is string manipulation, as we're constantly slicing and parsing substrings.
  • The algorithm uses recursion extensively to check all possible partitions, which inherently also uses a form of memoization due to overlapping subproblems within the recursive calls.
  • The overall approach is a backtracking problem, where we try out different substring splits and backtrack if we realize that the current path does not lead us to a solution.

This solution is efficient because it explores all possible partitions while avoiding unnecessary work, which is reduced by not exploring splits that would lead to values greater than the target sum x. As check progresses, it "cuts off" branches that cannot possibly result in a successful split, thereby streamlining the recursive search process.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Which two pointer techniques do you use to check if a string is a palindrome?

Example Walkthrough

Let's walk through an example to illustrate how the implemented solution works. Imagine we are tasked with finding the punishment number for n = 10. We want to check each number from 1 to 10 to see if its square can be split into substrings that sum up to itself.

Step 1: We start with i = 1.

  • We calculate its square: 1 * 1 = 1. There are no substrings to consider, so 1 is a valid partition by itself (1 = 1).
  • The check function would receive ("1", 0, 1) and return True, so we add 1 * 1 = 1 to the cumulative sum ans.

Step 2: Next, i = 2.

  • Square it: 2 * 2 = 4. Similar to 1, 4 does not have substrings to split, so it doesn't meet the condition.
  • The check function would receive ("4", 0, 2) and return False.

And we continue this process...

Step 3: For i = 3.

  • 3 * 3 = 9. There's no way to split 9 into substrings that add up to 3.
  • The check function would receive ("9", 0, 3) and return False.

Step 4: For i = 9 (Skipping a few steps for brevity).

  • 9 * 9 = 81. Here we can split 81 as 8 and 1, where 8 + 1 = 9.
  • The check function would receive ("81", 0, 9) and it would check the substring "8" (recursively calling check with "1", 1, 1) which returns True.
  • Therefore, we add 9 * 9 = 81 to the cumulative sum ans.

Continued Iteration: We would continue iterating over each number up to and including n, performing similar calculations.

Final Step:

  • After iterating from 1 through n (which is 10 in our example), we end up with ans holding the sum of valid squares, which, based on the provided numbers, are just 1 and 81.
  • The punishment number for n = 10 would therefore be 1 + 81 = 82.

This example demonstrates how the solution uses recursion to check possible splits for the square of each number and how the valid squares contribute to the punishment number.

Solution Implementation

1class Solution:
2    def punishmentNumber(self, n: int) -> int:
3        # Helper function to check if a square number can be split into two numbers that add up to the original number.
4        def can_split_to_original(num_str: str, start_index: int, target_sum: int) -> bool:
5            num_length = len(num_str)
6            # If the starting index is beyond the string's length and target sum is zero, this is valid.
7            if start_index >= num_length:
8                return target_sum == 0
9          
10            current_num = 0
11            # Consider different splits by incrementally including more digits.
12            for j in range(start_index, num_length):
13                current_num = current_num * 10 + int(num_str[j])
14                # If the current number exceeds the target, there's no point in continuing.
15                if current_num > target_sum:
16                    break
17                # Recursively check if the remaining string can fulfill the condition.
18                if can_split_to_original(num_str, j + 1, target_sum - current_num):
19                    return True
20            return False
21
22        # This variable holds the total sum of all square numbers fulfilling the condition.
23        total_sum = 0
24        # Iterate over the range 1 through n, inclusive.
25        for i in range(1, n + 1):
26            square_num = i * i  # Calculate the square of the number.
27            # If the squared number can be split as desired, add to total sum.
28            if can_split_to_original(str(square_num), 0, i):
29                total_sum += square_num
30      
31        # Return the total sum of all such square numbers.
32        return total_sum
33
1class Solution {
2    // Method to calculate the sum of special numbers up to n
3    public int punishmentNumber(int n) {
4        int sum = 0; // Initialize sum of special numbers
5        // Iterate through all numbers from 1 to n
6        for (int i = 1; i <= n; ++i) {
7            int square = i * i; // Calculate the square of the current number
8            // Check if the square can be expressed as a sum of numbers leading up to i
9            if (isSpecial(square, 0, i)) {
10                sum += square; // Add the square to the sum if it meets the condition
11            }
12        }
13        return sum; // Return the total sum of special numbers
14    }
15
16    // Helper method to check if a string representation of a number can be split
17    // into numbers adding up to a specific value
18    private boolean isSpecial(String numStr, int startIndex, int remaining) {
19        int length = numStr.length();
20        // Base case: if we've reached the end of the string
21        if (startIndex >= length) {
22            // Check if we've exactly reached the target sum
23            return remaining == 0;
24        }
25        int currentNumber = 0;
26        // Iterate through the substring starting at startIndex
27        for (int j = startIndex; j < length; ++j) {
28            // Add the next digit to the currentNumber
29            currentNumber = currentNumber * 10 + (numStr.charAt(j) - '0');
30            // If the currentNumber exceeds the remaining sum needed, break (optimization)
31            if (currentNumber > remaining) {
32                break;
33            }
34            // Recursive call to check the rest of the string with the updated remaining value
35            if (isSpecial(numStr, j + 1, remaining - currentNumber)) {
36                return true; // If successful, return true
37            }
38        }
39        return false; // If no combination adds up to remaining, return false
40    }
41}
42
1class Solution {
2public:
3    // Function to calculate the sum of special numbers
4    int punishmentNumber(int n) {
5        int sum = 0; // Initialize sum of the special numbers
6        // Loop through all numbers from 1 to n
7        for (int currentNum = 1; currentNum <= n; ++currentNum) {
8            int square = currentNum * currentNum; // Calculate the square of the current number
9            string squareStr = std::to_string(square); // Convert the square to a string
10            // Check if the square satisfies the condition
11            if (isSpecialNumber(squareStr, 0, currentNum)) {
12                sum += square; // Add the square to the sum if the condition is satisfied
13            }
14        }
15        return sum; // Return the sum of special numbers
16    }
17
18    // Helper function to check if the given string represents a special number
19    bool isSpecialNumber(const string& numStr, int index, int num) {
20        int strSize = numStr.size(); // Size of the number string
21        // If we've reached beyond the end of the string, check if num is 0 (completely decomposed)
22        if (index >= strSize) {
23            return num == 0;
24        }
25        int partialSum = 0; // Initialize the partial sum of digits from 'index'
26        // Iterate through the string starting from 'index'
27        for (int j = index; j < strSize; ++j) {
28            partialSum = partialSum * 10 + numStr[j] - '0'; // Build the number by appending the digit
29            // If the partial sum exceeds the remaining value of num, break (no further checks needed)
30            if (partialSum > num) {
31                break;
32            }
33            // Recurse to check if the remaining part of the string can make up the number num - partialSum
34            if (isSpecialNumber(numStr, j + 1, num - partialSum)) {
35                return true; // If the condition holds, return true
36            }
37        }
38        // If no valid decomposition found, return false
39        return false;
40    }
41};
42
1// Define the punishmentNumber function which calculates the sum of all the special numbers upto n.
2// Special numbers are those whose square can be split into two or more integers that sum to the original number.
3function punishmentNumber(n: number): number {
4    // Define helper function check to recursively determine if the square of a number
5    // can be split into two or more integers that sum to the original number.
6    const checkSplitSum = (squareAsString: string, index: number, targetSum: number): boolean => {
7        const length = squareAsString.length;
8        if (index >= length) {
9            return targetSum === 0;
10        }
11        let currentSum = 0;
12        for (let j = index; j < length; ++j) {
13            currentSum = currentSum * 10 + Number(squareAsString[j]);
14            if (currentSum > targetSum) {
15                break;
16            }
17            // Recursively check the next part of the string
18            if (checkSplitSum(squareAsString, j + 1, targetSum - currentSum)) {
19                return true;
20            }
21        }
22        return false;
23    };
24  
25    // Initialize the sum variable to keep track of the total sum of special numbers
26    let sumOfSpecialNumbers = 0;
27  
28    // Loop through each number from 1 to n to find all special numbers
29    for (let i = 1; i <= n; ++i) {
30        const square = i * i; // Calculate the square of the current number
31        const squareAsString = square.toString(); // Convert the square to a string for processing
32      
33        // Check if the square can be split into integers that sum to the original number
34        if (checkSplitSum(squareAsString, 0, i)) {
35            sumOfSpecialNumbers += square; // If it can be, add the square to the sum
36        }
37    }
38  
39    // Finally, return the sum of all special numbers found
40    return sumOfSpecialNumbers;
41}
42
Not Sure What to Study? Take the 2-min Quiz：

A person thinks of a number between 1 and 1000. You may ask any number questions to them, provided that the question can be answered with either "yes" or "no".

What is the minimum number of questions you needed to ask so that you are guaranteed to know the number that the person is thinking?

Time and Space Complexity

The given Python code defines a method punishmentNumber which calculates the sum of squares of certain integers up to n that satisfy a particular condition. The condition is verified using a helper method check that implements a recursive approach to determine if the square of the number can be represented as a consecutive sum of integers ending with the number itself.

Time Complexity:

The main time complexity comes from two sources:

  1. The outer for loop iterating from 1 to n.
  2. The recursive check function called for each i.

For each i, the check function explores different partitions of the digits of i*i into sums that could add up to i. In the worst case, this might result in checking each partition which leads to a search space that is exponential in the number of digits of i*i. However, the recursive search is pruned whenever a partition sum exceeds i, which can significantly reduce the number of searched partitions.

The number of digits d in i*i is O(log(i)) which means the worst case complexity of the check function is O(2^d). Since d is O(log(n)), the recursive calls for each i may reach a complexity of O(2^(2 * log(n))) due to the squaring operation, simplifying the expression to O(n^2) for each i. Multiplying this by n for the for-loop, we get a total worst-case time complexity of O(n^3).

Space Complexity:

The space complexity of the code is primarily determined by the maximum depth of the recursive call stack used by the check function.

The depth of the recursion is limited by the number of digits in x (the square of i), which is O(log(x)) = O(log(i^2)) = O(2*log(i)) = O(log(n)). Therefore, the space complexity is O(log(n)).

Note, the space used to store integers and for iteration is only O(1), which does not contribute significantly compared to the recursive stack.

Learn more about how to find time and space complexity quickly using problem constraints.

Fast Track Your Learning with Our Quick Skills Quiz:

Which two pointer techniques do you use to check if a string is a palindrome?

Recommended Readings

Math Basics

Math for Technical Interviews How much math do I need to know for technical interviews The short answer is about high school level math Computer science is often associated with math and some universities even place their computer science department under the math faculty However the reality is that you

Backtracking 1

Backtracking Template Prereq DFS with States problems dfs_with_states Combinatorial search problems Combinatorial search problems involve finding groupings and assignments of objects that satisfy certain conditions Finding all permutations combinations subsets and solving Sudoku are classic combinatorial problems The time complexity of combinatorial problems often grows rapidly with the size of

LeetCode Patterns - Top Patterns for Coding Interviews

Top Patterns to Conquer the Technical Coding Interview Should the written word bore you fear not A delightful video alternative awaits iframe width 560 height 315 src https www youtube com embed LW8Io6IPYHw title YouTube video player frameborder 0 allow accelerometer autoplay clipboard write encrypted media gyroscope picture in picture

Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.

←
↑TA 👨‍🏫