LeetCode 2698. Find the Punishment Number of an Integer Solution in Java, C++, Python & More | Explanation + Code

Description

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 in the range [1, 10] that satisfy the conditions in the statement:
- 1 since 1 * 1 = 1
- 9 since 9 * 9 = 81 and 81 can be partitioned into 8 and 1 with a sum equal to 8 + 1 == 9.
- 10 since 10 * 10 = 100 and 100 can be partitioned into 10 and 0 with a sum equal to 10 + 0 == 10.
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 in the range [1, 37] 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

Solutions

Solution 1: Enumeration + DFS

We enumerate i, where 1 ≤ i ≤ n. For each i, we split the decimal representation string of x = i², 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.

PythonJavaC++GoTypeScript

class Solution: def punishmentNumber(self, n: int) -> int: def check(s: str, i: int, x: int) -> bool: m = len(s) if i >= m: return x == 0 y = 0 for j in range(i, m): y = y * 10 + int(s[j]) if y > x: break if check(s, j + 1, x - y): return True return False ans = 0 for i in range(1, n + 1): x = i * i if check(str(x), 0, i): ans += x return ans(code-box)

class Solution { public int punishmentNumber(int n) { int ans = 0; for (int i = 1; i <= n; ++i) { int x = i * i; if (check(x + "", 0, i)) { ans += x; } } return ans; } private boolean check(String s, int i, int x) { int m = s.length(); if (i >= m) { return x == 0; } int y = 0; for (int j = i; j < m; ++j) { y = y * 10 + (s.charAt(j) - '0'); if (y > x) { break; } if (check(s, j + 1, x - y)) { return true; } } return false; } }(code-box)

class Solution { public: int punishmentNumber(int n) { int ans = 0; for (int i = 1; i <= n; ++i) { int x = i * i; string s = to_string(x); if (check(s, 0, i)) { ans += x; } } return ans; } bool check(const string& s, int i, int x) { int m = s.size(); if (i >= m) { return x == 0; } int y = 0; for (int j = i; j < m; ++j) { y = y * 10 + s[j] - '0'; if (y > x) { break; } if (check(s, j + 1, x - y)) { return true; } } return false; } };(code-box)

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 }(code-box)

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; }(code-box)