2698. Find the Punishment Number of an Integer - Solution

Q: Is Find the Punishment Number of an Integer easy or hard?

The problem is rated Medium because it requires recognizing a backtracking partition pattern. The implementation is short, but identifying that you must recursively split the digits of i² and prune invalid sums is the main challenge.

Q: Find the Punishment Number of an Integer Python/Java solution

In Python or Java, convert i² to a string and recursively try splitting the string into substrings. Convert each substring to an integer, accumulate the sum, and stop exploring when the sum exceeds i. If a partition ends exactly with sum equal to i, include i² in the total.

Q: How to solve Find the Punishment Number of an Integer in O(n)?

You can achieve near O(n) query time by precomputing which values satisfy the partition condition. Run the backtracking validation once for each i up to the maximum constraint and store the cumulative punishment sums. After preprocessing, answering a query for n only requires reading the stored prefix result.

Q: What is the best approach for Find the Punishment Number of an Integer?

The most practical approach uses backtracking to test whether the digits of i² can be partitioned into contiguous numbers that sum to i. For each i from 1 to n, recursively try different substring splits while tracking the current sum. The complexity is roughly O(n * 2^d), where d is the number of digits in i², which stays small in practice.

Q: Is Find the Punishment Number of an Integer asked at Google/Amazon/Meta?

Problems combining digit manipulation and backtracking appear frequently in interviews at companies like Amazon, Google, and Meta. While this exact problem may not always appear, the pattern of partitioning digits and validating sums is a common interview theme.

Q: What data structure is used in Find the Punishment Number of an Integer?

The core solution relies on recursion with backtracking and simple string or digit processing. Some implementations also use arrays to store prefix sums when applying precomputation. No complex data structures are required beyond recursion state tracking.

Q: What is the time complexity of Find the Punishment Number of an Integer?

The time complexity is O(n * 2^d). For each integer i from 1 to n, the algorithm explores possible partitions of the digits of i². Since the number of digits d is small (usually ≤7), the exponential partition exploration remains manageable.

The problem

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.

Example 1:

Input: n = 10
Output: 182
Explanation: There are exactly 3 integers i that satisfy the conditions in the statement:
- 1 since 1 * 1 = 1
- 9 since 9 * 9 = 81 and 81 can be partitioned into 8 + 1.
- 10 since 10 * 10 = 100 and 100 can be partitioned into 10 + 0.
Hence, the punishment number of 10 is 1 + 81 + 100 = 182

Example 2:

Input: n = 37
Output: 1478
Explanation: There are exactly 4 integers i that satisfy the conditions in the statement:
- 1 since 1 * 1 = 1. 
- 9 since 9 * 9 = 81 and 81 can be partitioned into 8 + 1. 
- 10 since 10 * 10 = 100 and 100 can be partitioned into 10 + 0. 
- 36 since 36 * 36 = 1296 and 1296 can be partitioned into 1 + 29 + 6.
Hence, the punishment number of 37 is 1 + 81 + 100 + 1296 = 1478

Constraints:

1 <= n <= 1000

Approach overview

Problem Overview: You need the punishment number for an integer n. For every i from 1 to n, square the number (i * i) and check whether its digits can be partitioned into contiguous pieces whose sum equals i. If it works, add i^2 to the total. The final sum of all such squares is the punishment number.

Approach 1: Backtracking for Partitioning Squares (O(n * 2^d) time, O(d) space)

Convert i * i into a string and try every possible way to split it into contiguous substrings. Use backtracking to explore partitions: at each position, extend the current segment and recursively check whether the running sum plus that segment can still reach i. If the sum equals i exactly when the string ends, the square qualifies. The search space is bounded by the number of digits d in i^2 (at most 7 for typical constraints), so even though the theoretical branching is 2^d, the recursion stays small. This approach is straightforward and mirrors the problem definition.

Approach 2: Precomputation and Dynamic Validation (O(n * 2^d) preprocessing, O(n) query, O(1) extra space)

Instead of recomputing the partition check every time, precompute which integers satisfy the property. For each i, run the same partition validation on the digits of i^2. If valid, store the square or maintain a running prefix sum of valid values. After preprocessing, computing the punishment number for any n becomes a simple prefix lookup or iteration. This approach combines math operations with recursive partition validation and works well when multiple queries are expected or when the platform constraints cap n to a small range.

Recommended for interviews: The backtracking partition approach is what interviewers expect. It shows you recognize the partition search pattern and can implement recursive pruning. Precomputation is an optimization layered on top of the same logic and demonstrates awareness of repeated computation and practical performance improvements.

Approach 1: Backtracking for Partitioning Squares

This approach involves recursively checking every possible partition of the number formed by the square of each integer i, up to n. The key observation is to find a way to partition the square such that the sum of its partitions equals the integer i itself.

Backtracking is used here to explore all potential partitions, checking if any valid partition matches the criteria. If found, we include the square of i in our sum to calculate the punishment number.

Python

def can_partition(square_str, i, sum_val):
    if sum_val > i:
        return False
    if not square_str:
        return sum_val == i
    num = 0
    for j in range(len(square_str)):
        num = num * 10 + int(square_str[j])
        if can_partition(square_str[j+1:], i, sum_val + num):
            return True
    return False


def punishment_number(n):
    punishment_sum = 0
    for i in range(1, n + 1):
        square = i * i
        if can_partition(str(square), i, 0):
            punishment_sum += square
    return punishment_sum


if __name__ == "__main__":
    n = 10
    print(f"Punishment Number of {n} = {punishment_number(n)}")

Java

import java.util.*;

public class PunishmentNumber {
    private static boolean canPartition(String squareStr, int i, int sum) {
        if (sum > i) return false;
        if (squareStr.isEmpty()) return sum == i;
        int num = 0;
        for (int j = 0; j < squareStr.length(); ++j) {
            num = num * 10 + (squareStr.charAt(j) - '0');
            if (canPartition(squareStr.substring(j + 1), i, sum + num))
                return true;
        }
        return false;
    }

    public static int punishmentNumber(int n) {
        int punishmentSum = 0;
        for (int i = 1; i <= n; ++i) {
            int square = i * i;
            String squareStr = String.valueOf(square);
            if (canPartition(squareStr, i, 0))
                punishmentSum += square;
        }
        return punishmentSum;
    }

    public static void main(String[] args) {
        int n = 10;
        System.out.println("Punishment Number of " + n + " = " + punishmentNumber(n));
    }
}

C++

#include <iostream>
#include <string>
using namespace std;

bool canPartition(const string& squareStr, int i, int sum) {
    if (sum > i) return false;
    if (squareStr.empty()) return sum == i;
    int num = 0;
    for (size_t j = 0; j < squareStr.size(); ++j) {
        num = num * 10 + (squareStr[j] - '0');
        if (canPartition(squareStr.substr(j + 1), i, sum + num))
            return true;
    }
    return false;
}

int punishmentNumber(int n) {
    int punishmentSum = 0;
    for (int i = 1; i <= n; ++i) {
        int square = i * i;
        string squareStr = to_string(square);
        if (canPartition(squareStr, i, 0))
            punishmentSum += square;
    }
    return punishmentSum;
}

int main() {
    int n = 10;
    cout << "Punishment Number of " << n << " = " << punishmentNumber(n) << endl;
    return 0;
}

JavaScript

function canPartition(squareStr, i, sum) {
    if (sum > i) return false;
    if (!squareStr) return sum === i;
    let num = 0;
    for (let j = 0; j < squareStr.length; ++j) {
        num = num * 10 + parseInt(squareStr[j]);
        if (canPartition(squareStr.slice(j + 1), i, sum + num))
            return true;
    }
    return false;
}

function punishmentNumber(n) {
    let punishmentSum = 0;
    for (let i = 1; i <= n; ++i) {
        let square = i * i;
        let squareStr = square.toString();
        if (canPartition(squareStr, i, 0))
            punishmentSum += square;
    }
    return punishmentSum;
}

const n = 10;
console.log(`Punishment Number of ${n} = ${punishmentNumber(n)}`);

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

// helper function to parse and partition
bool canPartition(const char* squareStr, int len, int i, int sum) {
    if (sum > i) return false; 
    if (len == 0) return sum == i;
    int num = 0;
    for (int j = 0; j<len; ++j) {
        num = num * 10 + (squareStr[j] - '0');
        if (canPartition(squareStr + j + 1, len - j - 1, i, sum + num))
            return true;
    }
    return false;
}

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

int main() {
    int n = 10;
    printf("Punishment Number of %d = %d\n", n, punishmentNumber(n));
    return 0;
}

using System;

class PunishmentNumber {
    private static bool CanPartition(string squareStr, int i, int sum) {
        if (sum > i) return false;
        if (squareStr.Length == 0) return sum == i;
        int num = 0;
        for (int j = 0; j < squareStr.Length; ++j) {
            num = num * 10 + (squareStr[j] - '0');
            if (CanPartition(squareStr.Substring(j + 1), i, sum + num))
                return true;
        }
        return false;
    }

    public static int GetPunishmentNumber(int n) {
        int punishmentSum = 0;
        for (int i = 1; i <= n; ++i) {
            int square = i * i;
            string squareStr = square.ToString();
            if (CanPartition(squareStr, i, 0))
                punishmentSum += square;
        }
        return punishmentSum;
    }

    static void Main() {
        int n = 10;
        Console.WriteLine("Punishment Number of " + n + " = " + GetPunishmentNumber(n));
    }
}

This Python solution involves a recursive function `can_partition` to check potential partitions of the squared value of each integer. It sums those squares which satisfy the criteria of splitting into parts whose sum matches the integer itself.

Complexity

Time Complexity: O(n * k), where n is the range till input integer and k is the partition-check based recursive operations.

Space Complexity: O(d), determined by recursive depth linked to square digits.

Try this approach in the editor

Approach 2: Precomputation and Dynamic Validation

In this approach, we try to optimize the solution by using precomputation and dynamic checking of the valid numbers whose square conforms to the partition rule. This method involves maintaining a list of valid squares till a certain number and using this as a lookup-table to quickly calculate the punishment number.

Python

def is_valid_partition(number, target):
    number_str = str(number)
    length = len(number_str)
    dp = [[False] * (target + 1) for _ in range(length + 1)]
    dp[0][0] = True

    for i in range(length):
        for t in range(target + 1):
            if dp[i][t]:
                num = 0
                for j in range(i, length):
                    num = num * 10 + int(number_str[j])
                    if t + num <= target:
                        dp[j + 1][t + num] = True

    return dp[length][target]


def punishment_number(n):
    punishment_sum = 0
    for i in range(1, n + 1):
        square = i * i
        if is_valid_partition(square, i):
            punishment_sum += square
    return punishment_sum


if __name__ == "__main__":
    n = 10
    print(f"Punishment Number of {n} = {punishment_number(n)}")

Java

import java.util.*;

public class DynamicPunishmentNumber {
    private static boolean isValidPartition(int number, int target) {
        String numberStr = String.valueOf(number);
        int len = numberStr.length();

        boolean[][] dp = new boolean[len + 1][target + 1];
        dp[0][0] = true;

        for (int i = 0; i < len; i++) {
            for (int t = 0; t <= target; t++) {
                if (dp[i][t]) {
                    int num = 0;
                    for (int j = i; j < len; j++) {
                        num = num * 10 + (numberStr.charAt(j) - '0');
                        if (t + num <= target) {
                            dp[j + 1][t + num] = true;
                        }
                    }
                }
            }
        }

        return dp[len][target];
    }

    public static int punishmentNumber(int n) {
        int punishmentSum = 0;
        for (int i = 1; i <= n; ++i) {
            int square = i * i;
            if (isValidPartition(square, i))
                punishmentSum += square;
        }
        return punishmentSum;
    }

    public static void main(String[] args) {
        int n = 10;
        System.out.println("Punishment Number of " + n + " = " + punishmentNumber(n));
    }
}

C++

#include <iostream>
#include <vector>
#include <string>
using namespace std;

bool is_valid_partition(int number, int target) {
    string numberStr = to_string(number);
    int len = numberStr.size();

    vector<vector<bool>> dp(len + 1, vector<bool>(target + 1, false));
    dp[0][0] = true;

    for (int i = 0; i < len; i++) {
        for (int t = 0; t <= target; t++) {
            if (dp[i][t]) {
                int num = 0;
                for (int j = i; j < len; j++) {
                    num = num * 10 + (numberStr[j] - '0');
                    if (t + num <= target) {
                        dp[j + 1][t + num] = true;
                    }
                }
            }
        }
    }

    return dp[len][target];
}

int punishmentNumber(int n) {
    int punishmentSum = 0;
    for (int i = 1; i <= n; ++i) {
        int square = i * i;
        if (is_valid_partition(square, i))
            punishmentSum += square;
    }
    return punishmentSum;
}

int main() {
    int n = 10;
    cout << "Punishment Number of " << n << " = " << punishmentNumber(n) << endl;
    return 0;
}

JavaScript

function isValidPartition(number, target) {
    const numberStr = number.toString();
    const len = numberStr.length;
    const dp = Array(len + 1).fill(0).map(() => Array(target + 1).fill(false));
    dp[0][0] = true;

    for (let i = 0; i < len; i++) {
        for (let t = 0; t <= target; t++) {
            if (dp[i][t]) {
                let num = 0;
                for (let j = i; j < len; j++) {
                    num = num * 10 + parseInt(numberStr[j]);
                    if (t + num <= target) {
                        dp[j + 1][t + num] = true;
                    }
                }
            }
        }
    }

    return dp[len][target];
}

function punishmentNumber(n) {
    let punishmentSum = 0;
    for (let i = 1; i <= n; ++i) {
        const square = i * i;
        if (isValidPartition(square, i))
            punishmentSum += square;
    }
    return punishmentSum;
}

const n = 10;
console.log(`Punishment Number of ${n} = ${punishmentNumber(n)}`);

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

bool is_valid_partition(int number, int target) {
    char numberStr[12];
    sprintf(numberStr, "%d", number);
    int len = strlen(numberStr);

    int dp[len + 1][target + 1];
    memset(dp, 0, sizeof(dp));
    dp[0][0] = 1;

    for (int i = 0; i < len; i++) {
        for (int t = 0; t <= target; t++) {
            if (dp[i][t]) {
                int num = 0;
                for (int j = i; j < len; j++) {
                    num = num * 10 + (numberStr[j] - '0');
                    if (t + num <= target) {
                        dp[j + 1][t + num] = 1;
                    }
                }
            }
        }
    }

    return dp[len][target] == 1;
}

int punishmentNumber(int n) {
    int punishmentSum = 0;
    for (int i = 1; i <= n; ++i) {
        int square = i * i;
        if (is_valid_partition(square, i))
            punishmentSum += square;
    }
    return punishmentSum;
}

int main() {
    int n = 10;
    printf("Punishment Number of %d = %d\n", n, punishmentNumber(n));
    return 0;
}

using System;

class DynamicPunishmentNumber {
    private static bool IsValidPartition(int number, int target) {
        string numberStr = number.ToString();
        int len = numberStr.Length;

        bool[,] dp = new bool[len + 1, target + 1];
        dp[0,0] = true;

        for (int i = 0; i < len; i++) {
            for (int t = 0; t <= target; t++) {
                if (dp[i,t]) {
                    int num = 0;
                    for (int j = i; j < len; j++) {
                        num = num * 10 + (numberStr[j] - '0');
                        if (t + num <= target) {
                            dp[j + 1,t + num] = true;
                        }
                    }
                }
            }
        }

        return dp[len,target];
    }

    public static int GetPunishmentNumber(int n) {
        int punishmentSum = 0;
        for (int i = 1; i <= n; ++i) {
            int square = i * i;
            if (IsValidPartition(square, i))
                punishmentSum += square;
        }
        return punishmentSum;
    }

    static void Main() {
        int n = 10;
        Console.WriteLine("Punishment Number of " + n + " = " + GetPunishmentNumber(n));
    }
}

Using a dp table, this Python solution determines whether the digit partitions of a squared integer can match the integer itself. The space and time efficiencies gained here come from dynamic pre-validation of sum consistencies.

Complexity

Time Complexity: O(n * m * target), with m as the number of digits and target being i.

Space Complexity: O(len * target), tied to dp table based on digit length and sum target.

Try this approach in the editor

Approach 3: Enumeration + DFS

We enumerate i, where 1 leq i leq n. For each i, we split the decimal representation string of x = i^2, and then check whether it meets the requirements of the problem. If it does, we add x to the answer.

After the enumeration ends, we return the answer.

The time complexity is O(n^{1 + 2 log_{10}^2}), and the space complexity is O(log n), where n is the given positive integer.

TypeScript

function punishmentNumber(n: number): number {
    const check = (s: string, i: number, x: number): boolean => {
        const m = s.length;
        if (i >= m) {
            return x === 0;
        }
        let y = 0;
        for (let j = i; j < m; ++j) {
            y = y * 10 + Number(s[j]);
            if (y > x) {
                break;
            }
            if (check(s, j + 1, x - y)) {
                return true;
            }
        }
        return false;
    };
    let ans = 0;
    for (let i = 1; i <= n; ++i) {
        const x = i * i;
        const s = x.toString();
        if (check(s, 0, i)) {
            ans += x;
        }
    }
    return ans;
}

func punishmentNumber(n int) (ans int) {
	var check func(string, int, int) bool
	check = func(s string, i, x int) bool {
		m := len(s)
		if i >= m {
			return x == 0
		}
		y := 0
		for j := i; j < m; j++ {
			y = y*10 + int(s[j]-'0')
			if y > x {
				break
			}
			if check(s, j+1, x-y) {
				return true
			}
		}
		return false
	}
	for i := 1; i <= n; i++ {
		x := i * i
		s := strconv.Itoa(x)
		if check(s, 0, i) {
			ans += x
		}
	}
	return
}

Try this approach in the editor

Complexity comparison

Approach	Complexity
Backtracking for Partitioning Squares	Time Complexity: *O(n k), where n is the range till input integer and k is the partition-check based recursive operations. Space Complexity: O(d)**, determined by recursive depth linked to square digits.
Precomputation and Dynamic Validation	Time Complexity: *O(n m * target), with m as the number of digits and target being i. Space Complexity: O(len * target)**, tied to dp table based on digit length and sum target.
Enumeration + DFS	–

Approach

Complexity

Backtracking for Partitioning Squares

Time Complexity: O(n * k), where n is the range till input integer and k is the partition-check based recursive operations.

Space Complexity: O(d), determined by recursive depth linked to square digits.

Precomputation and Dynamic Validation

Time Complexity: O(n * m * target), with m as the number of digits and target being i.

Space Complexity: O(len * target), tied to dp table based on digit length and sum target.

Enumeration + DFS

–

Complexity analysis

Approach	Time	Space	When to Use
Backtracking for Partitioning Squares	O(n * 2^d)	O(d)	Single query scenario where you directly test each square's digit partitions
Precomputation and Dynamic Validation	O(n * 2^d) preprocessing, O(n) query	O(1)	Multiple queries or platforms where valid punishment numbers can be cached

Frequently asked questions

Is Find the Punishment Number of an Integer easy or hard?

The problem is rated Medium because it requires recognizing a backtracking partition pattern. The implementation is short, but identifying that you must recursively split the digits of i² and prune invalid sums is the main challenge.

Find the Punishment Number of an Integer Python/Java solution

In Python or Java, convert i² to a string and recursively try splitting the string into substrings. Convert each substring to an integer, accumulate the sum, and stop exploring when the sum exceeds i. If a partition ends exactly with sum equal to i, include i² in the total.

How to solve Find the Punishment Number of an Integer in O(n)?

You can achieve near O(n) query time by precomputing which values satisfy the partition condition. Run the backtracking validation once for each i up to the maximum constraint and store the cumulative punishment sums. After preprocessing, answering a query for n only requires reading the stored prefix result.

What is the best approach for Find the Punishment Number of an Integer?

The most practical approach uses backtracking to test whether the digits of i² can be partitioned into contiguous numbers that sum to i. For each i from 1 to n, recursively try different substring splits while tracking the current sum. The complexity is roughly O(n * 2^d), where d is the number of digits in i², which stays small in practice.

Is Find the Punishment Number of an Integer asked at Google/Amazon/Meta?

Problems combining digit manipulation and backtracking appear frequently in interviews at companies like Amazon, Google, and Meta. While this exact problem may not always appear, the pattern of partitioning digits and validating sums is a common interview theme.

What data structure is used in Find the Punishment Number of an Integer?

The core solution relies on recursion with backtracking and simple string or digit processing. Some implementations also use arrays to store prefix sums when applying precomputation. No complex data structures are required beyond recursion state tracking.

What is the time complexity of Find the Punishment Number of an Integer?

The time complexity is O(n * 2^d). For each integer i from 1 to n, the algorithm explores possible partitions of the digits of i². Since the number of digits d is small (usually ≤7), the exponential partition exploration remains manageable.

Find the Punishment Number of an Integer solution and explanation

The problem

Approach overview

Approach 1: Backtracking for Partitioning Squares

Approach 2: Precomputation and Dynamic Validation

Approach 3: Enumeration + DFS

Complexity comparison

Complexity analysis

Frequently asked questions

Comments

Find the Punishment Number of an Integer solution and explanation

The problem

Approach overview

Approach 1: Backtracking for Partitioning Squares

Approach 2: Precomputation and Dynamic Validation

Approach 3: Enumeration + DFS

Complexity comparison

Complexity analysis

Frequently asked questions

Comments

Related problems