Find the Punishment Number of an Integer - Practice Coding Problems

Find the Punishment Number of an Integer - Problem

Math Backtracking Medium

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.

Input & Output

Example 1 — Basic Case

$ Input: n = 10

› Output: 182

💡 Note: Numbers 1, 9, 10 satisfy the condition: 1²=1 ("1"→1), 9²=81 ("8"+"1"→9), 10²=100 ("10"+"0"→10 or "1"+"00" invalid). Sum: 1+81+100=182

Example 2 — Small Input

$ Input: n = 37

› Output: 1478

💡 Note: All valid numbers up to 37 whose squares can be partitioned to sum back to the original number

Constraints

1 ≤ n ≤ 1000

Visualization

Tap to expand

Understanding the Visualization

Input

Given n=10, check numbers 1 through 10

Process

For each i, check if i² digits can partition to sum to i

Output

Sum of all valid i² values

Key Takeaway

🎯 Key Insight: Use backtracking to explore all possible ways to partition the square's digits into substrings that sum to the original number

Asked in

G Google 15 M Microsoft 10

The key insight is to use backtracking to try all possible partitions of each number's square. For each integer i from 1 to n, calculate i² and check if its digits can be partitioned into contiguous substrings that sum back to i. Best approach uses memoization to cache partition results. Time: O(n × 2ᵈ), Space: O(d)

Common Approaches

Approach	Time	Space	Notes
✓ Brute Force with Backtracking	O(n × 2ᵈ)	O(d)	For each number 1 to n, check if its square can be partitioned to sum back to the original number
Optimized with Memoization	O(n × 2ᵈ)	O(n × d)	Same backtracking approach but cache results to avoid recalculating the same partitions

Brute Force with Backtracking — Algorithm Steps

Step 1: Iterate through each number i from 1 to n
Step 2: Calculate square = i * i and convert to string
Step 3: Use backtracking to try all partitions of the square string
Step 4: If any partition sums to i, add i² to result

Visualization

Tap to expand

Step-by-Step Walkthrough

Calculate Square

For i=9, square = 81

Try Partitions

Try "8"+"1"=9 ✓

Add to Sum

Valid, so add 81 to punishment number

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>

bool canPartition(const char* s, int target, int start, int len) {
    if (start == len) {
        return target == 0;
    }
    
    for (int end = start + 1; end <= len; end++) {
        char substring[20];
        strncpy(substring, s + start, end - start);
        substring[end - start] = '\0';
        
        if (substring[0] == '0' && strlen(substring) > 1) {
            continue;
        }
        
        int num = atoi(substring);
        if (canPartition(s, target - num, end, len)) {
            return true;
        }
    }
    return false;
}

int solution(int n) {
    int punishment = 0;
    for (int i = 1; i <= n; i++) {
        int square = i * i;
        char squareStr[20];
        sprintf(squareStr, "%d", square);
        int len = strlen(squareStr);
        
        if (canPartition(squareStr, i, 0, len)) {
            punishment += square;
        }
    }
    return punishment;
}

int main() {
    int n;
    scanf("%d", &n);
    int result = solution(n);
    printf("%d\n", result);
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n × 2ᵈ)

For each of n numbers, we try 2^d partitions where d is digits in i²

✓ Linear Growth

Space Complexity

O(d)

Recursion stack depth equals number of digits in i²

✓ Linear Space

8.5K Views

Medium Frequency

~25 min Avg. Time

245 Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

Brute Force with Backtracking — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler