1150MediumMathBacktracking

Adding Up the Balanced Squares

Tracked in this browser only

Trains the technique from

LeetCode 2698Find the Punishment Number of an Integer

This is an original problem, written from a brief that listed the technique, the difficulty, the topics, the function shape and the input bounds — none of that problem's wording, examples, hints or editorials. The link is there so you can map your practice onto the standard set.

Same function shape, different story and different numbers.

A whole number i is balanced when the decimal digits of i * i can be cut into consecutive blocks whose values add up to i. For instance 9 is balanced, because 81 cuts into 8 and 1, which add to 9.

Return the total of i * i taken over every balanced i from 1 to n.

Examples

Example 1

Input: n = 9
Output: 82

Below ten only 1 and 9 are balanced: 1 is its own square, and 81 cuts into 8 and 1. Their squares add to 82.

Example 2

Input: n = 2
Output: 1

The square 4 cannot be cut into blocks adding to 2, so only 1 counts.

Example 3

Input: n = 45
Output: 3503

The balanced numbers up to 45 are 1, 9, 10, 36 and 45. The last of them counts because 2025 cuts into 20 and 25.

Constraints

1 <= n <= 1000

The signature

The editor is preloaded with this. It matches the parent problem's shape, so a solution that works here transfers to a judge unchanged.

Python

def punishment_number(n: int) -> int:

Java

public int punishmentNumber(int n)

The intended approach

Time O(n * 2^d) over the numbers and the cuts of their d-digit squares.Space O(d) for the recursion.

Each number stands on its own, so the work is one test repeated: can the digits of the square be cut into blocks adding to the number itself. Answer that by walking the digits from the front, choosing how long the next block is. At each position, grow a block one digit at a time and try the rest of the digits against the amount still owed. Blocks are consecutive and every digit belongs to exactly one block, so a cut is complete only when the digits run out with nothing left owing. The blocks all count positively, so a block already worth more than the amount owed can never be part of a cut, and neither can any longer block starting at the same place. Stopping there keeps the search tiny, since a square of `n` up to a thousand has at most seven digits.

Pythonentry point: punishment_number

def punishment_number(n):
    def splits(text, at, owed):
        if at == len(text):
            return owed == 0
        value = 0
        for end in range(at, len(text)):
            value = value * 10 + (ord(text[end]) - 48)
            if value > owed:
                break
            if splits(text, end + 1, owed - value):
                return True
        return False

    total = 0
    for i in range(1, n + 1):
        if splits(str(i * i), 0, i):
            total += i * i
    return total

The obvious version

For each number, list every way of cutting its square's digits by treating the gaps between digits as a bit pattern, build the blocks each pattern describes, add them up and see whether they reach the number. Costs two to the power of the digit gaps for every number.

This is not filler. It was written independently of the fast solution and the two were run against thousands of random inputs to confirm they agree — which is how the answer above is known to be right. Read it when the intended approach feels like a rabbit pulled out of a hat.

Python

def bf_punishment_number(n):
    total = 0
    for i in range(1, n + 1):
        text = str(i * i)
        width = len(text)
        ok = False
        for mask in range(1 << (width - 1)):
            pieces = []
            start = 0
            for cut in range(width - 1):
                if mask >> cut & 1:
                    pieces.append(text[start:cut + 1])
                    start = cut + 1
            pieces.append(text[start:])
            got = 0
            for piece in pieces:
                got += int(piece)
            if got == i:
                ok = True
                break
        if ok:
            total += i * i
    return total

How this problem was checked

Reference solution executed against 15 test cases — 15 passed
3 worked examples confirmed by running the code
0 eight-word sequences shared with the parent problem's text
Agrees with an independent brute-force implementation on random inputs
Test set proven to catch deliberately broken solutions
Reviewed by a separate quality pass
Function arity matches the parent problem's

Create a free account to unlock the solution

See the step-by-step animation, the intuition, and clean code in every language — free, no credit card.

Create a free account Already have an account? Sign in