2698. 求一个整数的惩罚数 - Find the Punishment Number of an Integer

🎯 问题描述（来源于LeetCode）

描述：
给定一个正整数 n，返回 1 到 n 之间所有惩罚数的平方和。
惩罚数的定义为：若一个整数 i 的平方 i*i 的十进制表示可以分割成若干连续子串，使得这些子串对应的整数值之和等于 i，则 i 是惩罚数。

说明：

分割时子串必须连续，且必须使用所有数字。
例如 1 是惩罚数，因为 1²=1 可分割为 "1"，和为 1。
例如 9 是惩罚数，因为 9²=81 可分割为 "8"+"1"，和为 9。

示例：

示例 1：

1
2
3

输入: n = 10
输出: 182
解释: 惩罚数为 1 和 9，平方和为 1+81=82？实际 LeetCode 示例输出为 82，请以题目为准。

示例 2：

1
2
3

输入: n = 37
输出: 1478
解释: 惩罚数有 1, 9, 10, 36，平方和为 1+81+100+1296=1478。

💻 解题思路

思路：DFS 通过参数模拟分割决策

算法步骤

定义辅助函数 can_partition(num)，判断 num 是否为惩罚数：
- 将 num * num 转为字符串 s，长度为 N。
- 定义递归函数 dfs(gap, total, cur)，其中：
  - gap 表示当前已处理的最后一个字符索引（从 0 开始）。
  - total 表示已经确定并累加的部分的和（不包括当前正在构建的数字 cur）。
  - cur 表示当前正在构建的数字（从当前分割的起始字符累加而成）。
- 递归终止：如果 gap == N - 1（已处理到最后一个字符），则判断 total + cur == num。
- 否则，取下一个字符 next_char = int(s[gap+1])，尝试两种操作：
  - 不分割：将下一个字符追加到当前数字 cur 上，递归 dfs(gap+1, total, cur * 10 + next_char)。
  - 分割：将当前数字 cur 加入 total，并以下一个字符作为新 cur 开始，递归 dfs(gap+1, total + cur, next_char)。
- 返回是否找到有效分割。
若字符串 s 为空（即 num=0），直接返回 False。
从 dfs(0, 0, int(s[0])) 开始搜索。
主函数遍历 i=1 到 n，若 can_partition(i) 为真，则累加 i*i。
返回累加结果。

代码实现

from typing import List

class Solution:
    def punishmentNumber(self, n: int) -> int:
        def can_partition(num: int) -> bool:
            s = str(num * num)
            N = len(s)

            def dfs(gap, total, cur):
                if gap == N - 1:
                    return total + cur == num
                next_char = int(s[gap+1])
                # 不分割：将下一个字符追加到当前数字
                if dfs(gap+1, total, cur * 10 + next_char):
                    return True
                # 分割：将当前数字加入总和，下一个字符作为新数字
                if dfs(gap+1, total + cur, next_char):
                    return True
                return False

            if not s:
                return False
            return dfs(0, 0, int(s[0]))
        
        ans = 0
        for i in range(1, n+1):
            if can_partition(i):
                ans += i * i
        return ans