Problem Description

You need to find the "punishment number" of a given positive integer n.

The punishment number is calculated by finding all integers i from 1 to n that satisfy a special property, squaring those integers, and summing up their squares.

The special property is: when you square the integer i to get i * i, you should be able to split the decimal digits of i * i into one or more contiguous groups such that when you add up the numeric values of these groups, the sum equals i.

For example, if i = 9:

• i * i = 81
• You can split "81" as "8" and "1"
• 8 + 1 = 9, which equals i
• So 9 satisfies the property and 81 would be included in the punishment number

Another example, if i = 10:

• i * i = 100
• You can split "100" as "10" and "0"
• 10 + 0 = 10, which equals i
• So 10 satisfies the property and 100 would be included in the punishment number

The solution uses a recursive helper function check to try all possible ways of partitioning the string representation of i * i. For each valid partition that sums to i, it adds i * i to the final answer.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• No: The problem involves checking integers and their squares, not nodes and edges as in graph problems.

Need to solve for kth smallest/largest?

• No: We're not finding the kth smallest or largest element, but rather summing squares that meet a specific condition.

Involves Linked Lists?

• No: The problem works with integers and their string representations, not linked lists.

Does the problem have small constraints?

• Yes: For each integer i, we need to check all possible ways to partition the string representation of i * i. The number of possible partitions grows exponentially with the number of digits, but since we're dealing with squares of integers up to n (typically n ≤ 1000), the string lengths are manageable (at most 7 digits for 1000² = 1,000,000).

Brute force / Backtracking?

• Yes: We need to explore all possible ways to partition each square's string representation and check if any partition sums to the original number. This requires backtracking to try different partition points and backtrack when a partition doesn't lead to the desired sum.

Conclusion: The flowchart correctly leads us to use a Backtracking approach. For each integer from 1 to n, we use backtracking to recursively try all possible ways to split the string representation of its square, checking if any valid partition sums to the original integer.

Intuition

The key insight is recognizing that we need to explore all possible ways to split a number's digits into groups. When we see "all possible ways," this immediately suggests a backtracking approach.

Think about what we're actually doing: for a number like 81 (which is 9²), we need to check if we can split "81" into contiguous parts that sum to 9. We could split it as:

• "81" (one group: 81)
• "8" + "1" (two groups: 8 + 1 = 9) ✓

For a longer number like 1296 (which is 36²), we have many more possibilities:

• "1296" (one group)
• "1" + "296" (two groups)
• "12" + "96" (two groups)
• "129" + "6" (two groups)
• "1" + "2" + "96" (three groups)
• ... and so on

This is a classic decision tree problem. At each position in the string, we face a decision: should we continue extending the current number, or should we "cut" here and start a new number? This branching decision structure naturally leads to a recursive backtracking solution.

The backtracking works by:

1. Starting from the beginning of the string
2. Trying to form a number by taking 1 digit, 2 digits, 3 digits, etc.
3. For each formed number, subtract it from our target sum and recursively check if we can partition the remaining string to get the remaining sum
4. If we reach the end of the string and our remaining sum is exactly 0, we found a valid partition
5. If at any point our current number exceeds the remaining sum, we can prune that branch (optimization)

This approach ensures we explore all possible partitions systematically without missing any valid combinations or checking the same partition twice.

Learn more about Math and Backtracking patterns.

Solution Approach

The solution implements the backtracking strategy through enumeration and DFS (Depth-First Search).

Main Loop - Enumeration: The outer structure iterates through all integers from 1 to n. For each integer i:

1. Calculate its square: x = i * i
2. Convert x to a string representation
3. Use the check function to determine if this square can be partitioned to sum to i
4. If valid, add x to the running answer

DFS Helper Function - check(s, i, x): This recursive function explores all possible partitions of the string s starting from index i, trying to reach the target sum x.

Parameters:

• s: The string representation of the squared number
• i: Current starting index in the string
• x: The remaining target sum we need to achieve

The DFS works as follows:

1. Base Case: If we've processed the entire string (i >= m), we check if the remaining target x equals 0. If yes, we found a valid partition.

2. Recursive Case: Starting from index i, try forming numbers of different lengths:

   • Initialize y = 0 to build the current number
   • For each position j from i to the end of string:
     • Extend the current number: y = y * 10 + int(s[j])
     • Pruning: If y > x, we can stop early since adding more digits will only make it larger
     • Recursively check if we can partition the rest: check(s, j + 1, x - y)
     • If the recursive call returns True, we found a valid partition

Example Walkthrough: For i = 9, x = 81, string s = "81":

• Start at index 0, target sum = 9
• Try "8" (y = 8): Recursively check remaining "1" with target 9 - 8 = 1
  • At index 1, try "1" (y = 1): Target is 1 - 1 = 0
  • Reached end with target = 0, return True
• Since we found a valid partition, add 81 to the answer

The time complexity is O(n × 2^d) where d is the maximum number of digits in any square, as we potentially explore all partitions for each number.

Example Walkthrough

Let's walk through finding the punishment number for n = 12.

We need to check each integer from 1 to 12 to see if it satisfies the special property.

Checking i = 1:

• Square: 1² = 1, string = "1"
• Can we partition "1" to sum to 1?
• Only partition: "1" = 1 ✓
• Add 1 to answer

Checking i = 9:

• Square: 9² = 81, string = "81"
• Can we partition "81" to sum to 9?
• Call check("81", 0, 9):
  • Try taking "8": y = 8, remaining target = 9 - 8 = 1
    • Call check("81", 1, 1):
      • Try taking "1": y = 1, remaining target = 1 - 1 = 0
      • Reached end with target = 0 → return True
  • Found valid partition "8" + "1" = 9 ✓
• Add 81 to answer

Checking i = 10:

• Square: 10² = 100, string = "100"
• Can we partition "100" to sum to 10?
• Call check("100", 0, 10):
  • Try taking "1": y = 1, remaining target = 10 - 1 = 9
    • Call check("100", 1, 9):
      • Try taking "0": y = 0, remaining = 9
        • Call check("100", 2, 9):
          • Try "0": y = 0, remaining = 9
          • Reached end with target = 9 ≠ 0 → return False
      • Try taking "00": y = 0, remaining = 9
        • Reached end with target = 9 ≠ 0 → return False
  • Try taking "10": y = 10, remaining target = 10 - 10 = 0
    • Call check("100", 2, 0):
      • Try taking "0": y = 0, remaining = 0 - 0 = 0
      • Reached end with target = 0 → return True
  • Found valid partition "10" + "0" = 10 ✓
• Add 100 to answer

Checking i = 11:

• Square: 11² = 121, string = "121"
• Can we partition "121" to sum to 11?
• Call check("121", 0, 11):
  • Try "1": y = 1, remaining = 10
    • Recursively check "21" for sum 10 → returns False
  • Try "12": y = 12, but 12 > 11 (pruned)
  • No valid partition found ✗
• Don't add to answer

Checking i = 12:

• Square: 12² = 144, string = "144"
• Can we partition "144" to sum to 12?
• Try all partitions:
  • "1" + "44" = 45 ✗
  • "1" + "4" + "4" = 9 ✗
  • "14" + "4" = 18 ✗
  • "144" = 144 ✗
• No valid partition found ✗

Final Answer: 1 + 81 + 100 = 182

The DFS efficiently explores the partition tree, pruning branches where the current number exceeds the remaining target, ensuring we find all valid partitions without redundant computation.

Solution Implementation

Time and Space Complexity

The time complexity is O(n^(1 + 2*log₁₀(2))) or approximately O(n^1.602).

Let me break down the analysis:

1. The outer loop runs from 1 to n, contributing O(n) iterations.

2. For each number i, we compute i² and convert it to a string. The square i² can have up to 2*log₁₀(i) digits, which is at most 2*log₁₀(n) digits.

3. The check function explores all possible ways to partition the string representation of i². In the worst case, for a string of length m, the number of recursive calls can be exponential in m. Specifically:

   • At each position, we can choose to split at any of the remaining positions
   • This creates a branching factor that leads to O(2^m) possible partitions in the worst case
   • Since m = O(log(i²)) = O(log(n²)) = O(2*log₁₀(n)), the check function has complexity O(2^(2*log₁₀(n)))

4. Using the property that 2^(log₁₀(n)) = n^(log₁₀(2)), we get:

   • O(2^(2*log₁₀(n))) = O((2^(log₁₀(n)))²) = O((n^(log₁₀(2)))²) = O(n^(2*log₁₀(2)))

5. The total time complexity is: O(n) * O(n^(2*log₁₀(2))) = O(n^(1 + 2*log₁₀(2)))

Since log₁₀(2) ≈ 0.301, this gives us O(n^1.602).

The space complexity is O(log n) due to:

• The recursion depth of the check function, which is bounded by the length of the string (at most O(log n) digits)
• The string storage for str(x) which takes O(log n) space

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

Common Pitfalls

1. Integer Overflow in Number Construction

When building current_number by repeatedly multiplying by 10 and adding digits, large sequences of digits can cause integer overflow in languages with fixed integer sizes. While Python handles arbitrary precision integers automatically, this could be problematic in other languages.

Solution: Add a check to ensure current_number doesn't exceed reasonable bounds, or use appropriate data types for your language.

2. Missing Early Termination Optimization

A critical optimization that's often overlooked is stopping the partition attempt when current_number > target_sum. Without this pruning, the algorithm explores unnecessary branches where the current number alone already exceeds what we need.

Example: For string "1000" with target 10, without pruning we'd try "100", "1000", etc., even though "100" already exceeds 10.

Solution: Always include the pruning condition:

if current_number > target_sum:
    break

3. Incorrect Base Case Handling

A common mistake is checking if start_idx == string_length instead of if start_idx >= string_length. This can cause issues if the recursion somehow skips indices or if the logic is modified later.

Solution: Use >= for defensive programming:

if start_idx >= string_length:
    return target_sum == 0

4. Leading Zeros in Partitions

The current solution treats "01" as 1, which might not be the intended behavior. For example, if we have "100" and partition it as "1", "00", the "00" becomes 0, which works but might violate the problem's intent if leading zeros aren't allowed in multi-digit numbers.

Solution: If leading zeros should invalidate a partition, add a check:

# Skip partitions with leading zeros (except single zero)
if num_str[start_idx] == '0' and end_idx > start_idx:
    continue

5. Forgetting to Convert Square to String

Attempting to partition the integer directly instead of its string representation will cause errors. The algorithm relies on string indexing and character-by-character processing.

Solution: Always convert to string before calling the partition function:

if can_partition_to_target(str(square), 0, num):  # Don't forget str()