Find the Punishment Number of an Integer
Problem
Given a positive integer n, return its punishment number: the sum of i*i over every integer i with 1 ≤ i ≤ n such that the decimal digits of i*i can be split into contiguous substrings whose integer values add up to exactly i.
n = 10182def punishmentNumber(n):
# Can digits of s[start:] be split so the parts sum to target?
def can(s, start, target):
if start == len(s):
return target == 0 # used all digits, hit target exactly
cur = 0
for end in range(start, len(s)):
cur = cur * 10 + int(s[end]) # extend current substring
if cur > target:
break # too big, and only grows -> prune
if can(s, end + 1, target - cur):
return True # this cut works
return False
total = 0
for i in range(1, n + 1):
if can(str(i * i), 0, i): # square can be partitioned into i
total += i * i
return totalfunction punishmentNumber(n) {
// Can digits of s from `start` be split so parts sum to target?
function can(s, start, target) {
if (start === s.length) return target === 0;
let cur = 0;
for (let end = start; end < s.length; end++) {
cur = cur * 10 + (s.charCodeAt(end) - 48); // extend substring
if (cur > target) break; // prune: only grows
if (can(s, end + 1, target - cur)) return true;
}
return false;
}
let total = 0;
for (let i = 1; i <= n; i++) {
if (can(String(i * i), 0, i)) total += i * i;
}
return total;
}int punishmentNumber(int n) {
int total = 0;
for (int i = 1; i <= n; i++) {
String s = Integer.toString(i * i);
if (can(s, 0, i)) total += i * i;
}
return total;
}
// Can digits of s from `start` be split so parts sum to target?
boolean can(String s, int start, int target) {
if (start == s.length()) return target == 0;
int cur = 0;
for (int end = start; end < s.length(); end++) {
cur = cur * 10 + (s.charAt(end) - '0'); // extend substring
if (cur > target) break; // prune: only grows
if (can(s, end + 1, target - cur)) return true;
}
return false;
}bool can(const string& s, int start, int target) {
if (start == (int)s.size()) return target == 0;
int cur = 0;
for (int end = start; end < (int)s.size(); end++) {
cur = cur * 10 + (s[end] - '0'); // extend substring
if (cur > target) break; // prune: only grows
if (can(s, end + 1, target - cur)) return true;
}
return false;
}
int punishmentNumber(int n) {
int total = 0;
for (int i = 1; i <= n; i++) {
if (can(to_string(i * i), 0, i)) total += i * i;
}
return total;
}Explanation
The punishment number adds up i*i for each i that passes a single test: can the digit string of i*i be cut into contiguous pieces that sum to i? Everything hinges on that yes/no question, and we answer it with backtracking.
Think of the digits of i*i laid out in a row. A partition is a choice of cut positions between adjacent digits. The helper can(s, start, target) tries every length of the first piece starting at start: it reads one more digit at a time into cur, then recurses on the remainder with the goal reduced to target - cur.
Two base/prune rules keep it fast. When start reaches the end of the string, the split is valid only if target == 0 — every digit was consumed and the running sum landed exactly on i. And because cur only grows as we read more digits, once cur > target we break: no longer prefix can help.
For example i = 9 gives 81. We try 8 as the first piece (target drops to 1), then 1 (target drops to 0) with no digits left — success, so 9 qualifies and we add 81. For i = 10, 100 splits as 10 + 0 = 10.
The squares are tiny (at most 7 digits for n = 1000), so the partition search is cheap. We simply loop i from 1 to n, run the check, and sum the squares that pass.