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

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

✓ Binary Search

⏱️ Time: N/A Space: N/A

Brute Force with Backtracking

⏱️ Time: O(n × 2ᵈ) Space: O(d)

For each integer i from 1 to n, calculate i². Then use backtracking to try all possible ways to partition the digits of i² into contiguous substrings. If any partition sums to i, add i² to the punishment number.

Optimized with Memoization

⏱️ Time: O(n × 2ᵈ) Space: O(n × d)

Uses the same backtracking logic but adds memoization to cache whether a string can be partitioned to reach a target sum. This avoids redundant calculations when the same partition problem appears multiple times.

Algorithm Steps — Algorithm Steps

Code -

solution.c — C

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

int canPartition(const char* s, int target, int start, int len) {
    if (start == len) {
        return target == 0;
    }
    
    for (int end = start + 1; end <= len; end++) {
        if (s[start] == '0' && (end - start) > 1) {
            continue;
        }
        int num = 0;
        for (int i = start; i < end; i++) {
            num = num * 10 + (s[i] - '0');
        }
        if (canPartition(s, target - num, end, len)) {
            return 1;
        }
    }
    return 0;
}

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() {
    char line[100];
    fgets(line, sizeof(line), stdin);
    
    int n = 0;
    for (int i = 0; line[i]; i++) {
        if (line[i] >= '0' && line[i] <= '9') {
            n = n * 10 + (line[i] - '0');
        }
    }
    
    int result = solution(n);
    printf("%d\n", result);
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

✓ Linear Growth

Space Complexity

⚡ Linearithmic 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