Find the Punishment Number of an Integer

Medium

Asked by:

21 views

Topics:

RecursionDynamic Programming

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

Try coding it in LeetCode

Solution

Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

What is the upper bound for the integer 'n'?
Should I handle the case where n is zero or negative?
If a number's square cannot be partitioned into a sum of numbers equal to the original number, should I include it in the punishment number calculation?
Can you clarify the definition of 'partition'? Can the square of the number be split into any number of parts, and do these parts need to be contiguous?
If `n` is very large, and the punishment number is also expected to be large, what is the required data type to return the result? Should I use `long` or `BigInteger`?

Brute Force Solution

Approach

To find the punishment number, we need to check every number from 1 up to the input number. For each number, we square it, then try all possible ways to split the squared number into separate parts, and finally, we check if the sum of these parts equals the original number.

Here's how the algorithm would work step-by-step:

Start with the number 1 and go up to the input number.
For each of these numbers, calculate its square.
Now, imagine the square as a string of digits. We're going to explore every possible way to cut this string into different pieces.
For each way of cutting the string, add up the numbers that each piece represents.
If the sum of the numbers from the pieces equals the original number we started with, then that number is a valid 'punishment number'.
If a number is a valid punishment number, we remember its square.
After checking all numbers up to the input, add up all the squares we remembered. That sum is the final answer.

Code Implementation

def find_punishment_number(number_limit):
    total_punishment_number = 0

    for current_number in range(1, number_limit + 1):
        square_of_number = current_number * current_number
        square_as_string = str(square_of_number)

        if is_punishment_number(square_as_string, current_number, 0, 0):
            total_punishment_number += square_of_number

    return total_punishment_number

def is_punishment_number(squared_string, original_number, current_index, current_sum):
    # Base case: If we've reached the end of the squared string
    if current_index == len(squared_string):
        return current_sum == original_number

    for i in range(current_index, len(squared_string)): 
        # Iterate through possible split points.
        prefix = squared_string[current_index:i + 1]
        prefix_number = int(prefix)

        # Recursively check the remaining string with the updated sum
        if is_punishment_number(squared_string, original_number, i + 1, current_sum + prefix_number):
            return True

    # If no combination works, return False
    return False

Big(O) Analysis

Time Complexity

O(n * 4^k) – The algorithm iterates from 1 to n. Inside the loop, the square of the number is computed (O(1)). The core complexity lies in splitting the squared number's digits. If the squared number has k digits, there are approximately 2^(k-1) ways to split it (treating it as a string), although the problem statement refers to roughly 4^k for the split checking complexity. Each split involves summing the resulting numbers, which takes O(k) time. Hence, the splitting and checking process takes O(4^k) time, where k is the number of digits in the square, which is at most 2*log10(n). Since the outer loop iterates n times, the overall time complexity is approximately O(n * 4^k).

Space Complexity

O(log N) – The primary driver of space complexity is the recursion depth of the splitting process within the isPunishmentNumber function. The depth of the recursion is proportional to the number of digits in the square of the input number i. Since i ranges from 1 to N, the maximum square is N*N. The number of digits in N*N is proportional to log(N*N) = 2*log(N), which simplifies to log N. Therefore the auxiliary space used by the call stack is O(log N).

Optimal Solution

Approach

The punishment number is calculated by summing the punishment numbers of each integer up to a given input. To find the punishment number for a single integer, we square the integer and then check if that square can be split into parts such that the sum of the parts equals the original integer. We use a clever checking method to avoid trying every possible split.

Here's how the algorithm would work step-by-step:

Start with the given number and work backward to 1.
For each number, square it.
Now, check if the squared number can be split into pieces that add up to the original number. We do this by trying different ways to cut the squared number into parts.
The trick is to use a technique like 'depth-first search' to explore the cutting possibilities. Think of it as trying all the possible ways to break the square into sections.
When splitting, keep track of the sum of the pieces you've cut so far. Also, make sure the sum never exceeds the original number, and make sure the pieces are formed by concatenating consecutive digits.
If you find a way to split the square so that the pieces add up to the original number, then the square is a 'punishment' and keep it.
Once you have checked all the numbers from the input down to 1, add up all the 'punishment' squares you found. This total is the 'punishment number' for the given input.

Code Implementation

def find_punishment_number(maximum_number):
    total_punishment_number = 0
    for number in range(1, maximum_number + 1):
        squared_number = number * number
        squared_number_string = str(squared_number)

        # Check if the squared number can be partitioned to sum to the original number.
        if can_partition_into_sum(squared_number_string, number, 0, 0):
            total_punishment_number += squared_number

    return total_punishment_number

def can_partition_into_sum(number_string, target_number, current_index, current_sum):
    if current_index == len(number_string):
        return current_sum == target_number

    for i in range(current_index, len(number_string)): 
        # Iterate over all possible ending positions for current section.
        section_string = number_string[current_index:i + 1]
        section_value = int(section_string)
        
        #The check avoids unnecessary calculations
        if current_sum + section_value <= target_number:

            #Use DFS to see if the remaining number can be partitioned.
            if can_partition_into_sum(number_string, target_number, i + 1, current_sum + section_value):
                return True

    return False

Big(O) Analysis

Time Complexity

O(n * 4^len(n^2)) – The algorithm iterates from n down to 1, so the outer loop contributes O(n). Inside the loop, we square the number, which is O(1) as it's a fixed operation. The core complexity comes from checking if the square can be split into parts summing to the original number using a recursive depth-first search. The number of possible splits of the squared number (whose length is at most 2*len(n)) is exponential, specifically O(4^len(n^2)) since each digit can potentially start a new number in the partition. Therefore the overall time complexity is O(n * 4^len(n^2)).

Space Complexity

O(log N) – The space complexity is dominated by the depth-first search (DFS) recursion. The maximum depth of the recursion is determined by the number of digits in the square of the input number N, which is proportional to log(N^2) or 2*log(N). Since the base of the logarithm is a constant, we can simplify this to log N. Each recursive call uses a constant amount of space for local variables and function call overhead. Therefore, the auxiliary space used by the recursion stack is O(log N).

Edge Cases

Case	How to Handle
Input n is 0	Return 0 since the punishment number is defined only for positive integers.
Input n is 1	Return 1 directly as 1^2 = 1 and 1 is a valid split.
Integer overflow when calculating i*i	Use a long data type to store the result of i*i before passing to the splitting function to avoid overflow.
n is a large value close to the upper bound	The algorithm's runtime will scale with the value of n; consider optimization if performance becomes a bottleneck, but a brute force approach is acceptable for reasonable constraints.
No valid split exists for a given i*i	The recursive helper function will return false, indicating that the current i does not contribute to the punishment number.
Single-digit numbers squared	Handle the base case in the recursive helper function where the squared value is already a single digit and simply matches the expected remaining sum.
Input n is a negative number	Return 0 or throw an IllegalArgumentException, as the problem specifies positive integers.
Maximum possible integer input that results in maximum recursion depth due to many possible splits.	Ensure stack overflow does not occur due to excessive recursion, potentially by using iterative approach or setting reasonable recursion depth.

Input n is 0

How to Handle:

Return 0 since the punishment number is defined only for positive integers.

Input n is 1

How to Handle:

Return 1 directly as 1^2 = 1 and 1 is a valid split.

Integer overflow when calculating i*i

How to Handle:

Use a long data type to store the result of i*i before passing to the splitting function to avoid overflow.

n is a large value close to the upper bound

How to Handle:

The algorithm's runtime will scale with the value of n; consider optimization if performance becomes a bottleneck, but a brute force approach is acceptable for reasonable constraints.

No valid split exists for a given i*i

How to Handle:

The recursive helper function will return false, indicating that the current i does not contribute to the punishment number.

Single-digit numbers squared

How to Handle:

Handle the base case in the recursive helper function where the squared value is already a single digit and simply matches the expected remaining sum.

Input n is a negative number

How to Handle:

Return 0 or throw an IllegalArgumentException, as the problem specifies positive integers.

Maximum possible integer input that results in maximum recursion depth due to many possible splits.

How to Handle:

Ensure stack overflow does not occur due to excessive recursion, potentially by using iterative approach or setting reasonable recursion depth.