LeetCode 2698 - Find the Punishment Number of an Integer

The problem asks us to compute a special value called the punishment number for a given integer n. For every integer i in the range [1, n], we square the number and examine the decimal representation of i i.

leetcodemediummathbacktracking

LeetCode Problem 2698

Difficulty: 🟡 Medium
Topics: Math, Backtracking

Solution

Problem Understanding

The problem asks us to compute a special value called the punishment number for a given integer n.

For every integer i in the range [1, n], we square the number and examine the decimal representation of i * i. The key question is whether we can split this decimal string into contiguous pieces such that the sum of those pieces equals the original number i.

If a number satisfies this condition, then its square contributes to the final answer.

For example, consider i = 36.

36 * 36 = 1296

We can partition "1296" into:

1 + 29 + 6 = 36

Since the partition sum equals 36, the square 1296 is included in the punishment number.

The input consists of a single integer n, where:

1 <= n <= 1000

The output is a single integer representing the sum of all valid squares.

The constraint is relatively small. Even though checking every possible partition sounds exponential, the number of digits in i^2 is tiny because:

1000^2 = 1000000

This means the longest string we ever process has only 7 digits. That observation is extremely important because it makes backtracking feasible.

Several edge cases are important:

  • Single digit squares such as 1, where no partitioning is needed.
  • Numbers whose square contains zeros, such as 100, where partitions like "10" and "0" are valid.
  • Numbers where many partition combinations exist, which could cause excessive recursion in a naive implementation.
  • The smallest input n = 1.

Approaches

Brute Force Approach

The brute force idea is straightforward.

For every integer i from 1 to n:

  1. Compute square = i * i
  2. Convert the square into a string
  3. Generate every possible partition of the string
  4. Compute the sum of each partition
  5. Check whether any partition sum equals i

A string with length m has 2^(m-1) possible partition patterns because every gap between digits can either contain a cut or not contain a cut.

This approach is correct because it explicitly checks every valid partition configuration. If a valid partition exists, brute force will eventually find it.

However, generating all partitions explicitly can become inefficient and repetitive. Many recursive paths recompute the same intermediate states.

Optimal Approach

The better approach still uses backtracking, but performs the partition generation incrementally.

Instead of constructing all partitions first, we recursively build substrings while tracking the current accumulated sum.

At each position in the string:

  • Extend the current substring digit by digit
  • Convert it into a number
  • Add it to the running total
  • Recursively process the remaining suffix

The key optimization is pruning.

If the running sum already exceeds the target i, there is no reason to continue exploring that branch.

Since the square string length is at most 7, the recursive search remains very small.

Approach Time Complexity Space Complexity Notes
Brute Force O(n × 2^d × d) O(d) Generates all partition patterns explicitly
Optimal O(n × 2^d) O(d) Backtracking with pruning, where d is the number of digits in i²

Since d <= 7, the algorithm is effectively very fast in practice.

Algorithm Walkthrough

Step 1: Iterate Through Every Number

Loop through every integer i from 1 to n.

For each number:

square = i * i

Convert the square into a string so we can partition its digits.

We now need to determine whether the digit string can be partitioned into contiguous substrings whose sum equals i.

We use a recursive helper function with:

  • Current index in the string
  • Current accumulated sum
  • Target value i

Step 3: Build Substrings Incrementally

At each recursive call:

  1. Start from the current index
  2. Extend the substring one digit at a time
  3. Convert the substring into an integer
  4. Add it to the running sum
  5. Recurse on the remaining suffix

For example, for "1296":

1 | 296
12 | 96
129 | 6
1296

Each split point creates a new recursive branch.

Step 4: Prune Impossible Paths

If the running sum exceeds the target i, stop exploring that branch.

This pruning is important because all numbers are non-negative. Once we exceed the target, adding more numbers can never bring the sum back down.

Step 5: Check the Base Case

When we reach the end of the string:

  • If the accumulated sum equals i, return True
  • Otherwise return False

Step 6: Add Valid Squares

If the recursive search succeeds:

answer += i * i

Continue until all numbers from 1 to n have been checked.

Why it works

The recursive search explores every possible way to partition the square string into contiguous substrings. Every recursive path corresponds to exactly one partition configuration.

Because all partition possibilities are explored, if a valid decomposition exists, the algorithm will find it. If no recursive branch reaches the target sum, then no valid partition exists.

The pruning step is safe because all substring values are non-negative. Once the running total exceeds the target, no future additions can reduce it.

Python Solution

class Solution:
    def punishmentNumber(self, n: int) -> int:
        def can_partition(square_str: str, index: int, current_sum: int, target: int) -> bool:
            if index == len(square_str):
                return current_sum == target

            if current_sum > target:
                return False

            number = 0

            for end in range(index, len(square_str)):
                number = number * 10 + int(square_str[end])

                if can_partition(
                    square_str,
                    end + 1,
                    current_sum + number,
                    target
                ):
                    return True

            return False

        punishment_sum = 0

        for value in range(1, n + 1):
            square = value * value

            if can_partition(str(square), 0, 0, value):
                punishment_sum += square

        return punishment_sum

The implementation begins with the recursive helper function can_partition.

The helper receives:

  • The square string
  • The current index
  • The accumulated sum
  • The target value

The base case occurs when the index reaches the end of the string. At that point, we check whether the accumulated sum matches the target.

The variable number is built incrementally inside the loop. Instead of repeatedly slicing strings and converting them into integers, we construct the number digit by digit:

number = number * 10 + int(square_str[end])

This avoids unnecessary substring allocations.

For every possible substring ending position, we recursively continue searching from the next index.

If any recursive branch succeeds, we immediately return True.

The outer loop checks every integer from 1 to n. Whenever a valid partition exists, we add the square to the running punishment sum.

Go Solution

package main

import "strconv"

func punishmentNumber(n int) int {
	var canPartition func(string, int, int, int) bool

	canPartition = func(squareStr string, index int, currentSum int, target int) bool {
		if index == len(squareStr) {
			return currentSum == target
		}

		if currentSum > target {
			return false
		}

		number := 0

		for end := index; end < len(squareStr); end++ {
			number = number*10 + int(squareStr[end]-'0')

			if canPartition(squareStr, end+1, currentSum+number, target) {
				return true
			}
		}

		return false
	}

	punishmentSum := 0

	for value := 1; value <= n; value++ {
		square := value * value
		squareStr := strconv.Itoa(square)

		if canPartition(squareStr, 0, 0, value) {
			punishmentSum += square
		}
	}

	return punishmentSum
}

The Go implementation follows the same recursive structure as the Python version.

One notable difference is digit conversion. Instead of calling a parsing function repeatedly, the code converts digits directly using:

int(squareStr[end] - '0')

This is more efficient than repeatedly slicing strings and calling strconv.Atoi.

Go also requires explicitly declaring the recursive function variable before assigning the function body because the function references itself recursively.

Integer overflow is not a concern because the largest square is:

1000² = 1,000,000

which easily fits inside Go's int.

Worked Examples

Example 1

Input:

n = 10

We check every number from 1 to 10.

i i² Valid Partition Sum Included
1 1 1 1 Yes
2 4 4 4 No
3 9 9 9 No
4 16 1 + 6 7 No
5 25 2 + 5 7 No
6 36 3 + 6 9 No
7 49 4 + 9 13 No
8 64 6 + 4 10 No
9 81 8 + 1 9 Yes
10 100 10 + 0 10 Yes

The punishment number is:

1 + 81 + 100 = 182

Recursive Trace for 36

We test:

36² = 1296

Recursive exploration:

Current Partition Running Sum
1 1
1 + 2 3
1 + 29 30
1 + 29 + 6 36

Once the sum reaches 36 at the end of the string, the search succeeds.

Example 2

Input:

n = 37

Valid numbers are:

i i² Valid Partition
1 1 1
9 81 8 + 1
10 100 10 + 0
36 1296 1 + 29 + 6

Final answer:

1 + 81 + 100 + 1296 = 1478

Complexity Analysis

Measure Complexity Explanation
Time O(n × 2^d) Each square string explores all partition combinations
Space O(d) Recursive call stack depth equals digit count

Here, d is the number of digits in i².

Because:

n <= 1000

we have:

d <= 7

So the recursion tree is extremely small. Even though the theoretical complexity is exponential in digit count, the practical runtime is tiny.

Test Cases

solution = Solution()

assert solution.punishmentNumber(1) == 1  # Minimum input
assert solution.punishmentNumber(10) == 182  # Example 1
assert solution.punishmentNumber(37) == 1478  # Example 2

assert solution.punishmentNumber(2) == 1  # Only 1 qualifies
assert solution.punishmentNumber(9) == 82  # 1 + 81
assert solution.punishmentNumber(36) == 1478  # Includes 36^2

assert solution.punishmentNumber(8) == 1  # No additional valid numbers
assert solution.punishmentNumber(1000) > 0  # Upper constraint stress test

assert solution.punishmentNumber(45) == 3503  # Additional known valid range
Test Why
n = 1 Validates smallest possible input
n = 10 Verifies first official example
n = 37 Verifies second official example
n = 2 Ensures non-qualifying numbers are skipped
n = 9 Tests partition into two substrings
n = 36 Tests multi-part recursive partitioning
n = 8 Ensures invalid partitions are rejected
n = 1000 Stress test for upper constraint
n = 45 Tests larger cumulative result

Edge Cases

One important edge case is the smallest possible input, n = 1. In this case, the square is also 1, and no partitioning is required. Some implementations mistakenly assume at least one split must occur. This implementation correctly allows the entire string to remain as a single partition.

Another important edge case involves zeros inside the square representation. For example:

10² = 100

The valid partition is:

10 + 0

An incorrect implementation might reject partitions containing leading or standalone zeros. This solution treats every substring as a valid integer conversion, including "0".

A third edge case occurs when many recursive branches exist. For example, strings with repeated digits can create many partition combinations. Without pruning, recursion could explore unnecessary paths. The condition:

if current_sum > target:
    return False

prevents wasted exploration and keeps the recursion efficient.

Another subtle edge case occurs when no valid partition exists. The recursion must exhaustively explore every partition possibility before returning False. This implementation guarantees correctness because every possible split point is explored systematically.