求一个整数的惩罚数

标签: 数学 回溯

难度: Medium

给你一个正整数 n ，请你返回 n 的 惩罚数 。

n 的 惩罚数 定义为所有满足以下条件 i 的数的平方和：

  • 1 <= i <= n
  • i * i 的十进制表示的字符串可以分割成若干连续子字符串，且这些子字符串对应的整数值之和等于 i 。

示例 1：

输入：n = 10
输出：182
解释：总共有 3 个整数 i 满足要求：
- 1 ，因为 1 * 1 = 1
- 9 ，因为 9 * 9 = 81 ，且 81 可以分割成 8 + 1 。
- 10 ，因为 10 * 10 = 100 ，且 100 可以分割成 10 + 0 。
因此，10 的惩罚数为 1 + 81 + 100 = 182

示例 2：

输入：n = 37
输出：1478
解释：总共有 4 个整数 i 满足要求：
- 1 ，因为 1 * 1 = 1
- 9 ，因为 9 * 9 = 81 ，且 81 可以分割成 8 + 1 。
- 10 ，因为 10 * 10 = 100 ，且 100 可以分割成 10 + 0 。
- 36 ，因为 36 * 36 = 1296 ，且 1296 可以分割成 1 + 29 + 6 。
因此，37 的惩罚数为 1 + 81 + 100 + 1296 = 1478

提示：

  • 1 <= n <= 1000

Submission

运行时间: 46 ms

内存: 16.0 MB

from typing import *

ans = [0] * 1001
for n in range(1, 1001):
    sq = str(n**2)

    def dfs(start, s):
        if start == len(sq):
            return s == n
        res =False
        for i in range(start + 1, len(sq) + 1):
            res = dfs(i, s + int(sq[start:i])) or res
        return res

    ans[n] = n**2 if dfs(0, 0) else 0
for i in range(1, 1001):
    ans[i] += ans[i - 1]
# print(ans)


class Solution:
    def punishmentNumber(self, n: int) -> int:
        return ans[n]

Explain

首先，这个题解通过预计算的方法解决问题，即预先计算了所有可能的 n (从 1 到 1000) 的惩罚数，并将结果存储在数组 ans 中。对于每个 n，计算 n 的平方并转换为字符串 sq。然后，使用一个递归函数 dfs 来尝试所有可能的连续子字符串分割方式，查看这些子字符串的整数和是否等于 n。如果存在这样的分割方式，那么该 n 的惩罚数就是 n 的平方；否则为 0。最后，累加数组 ans 中的值，以便直接通过索引快速回答任何小于等于 1000 的 n 的询问。

时间复杂度: O(n * 2^k)

空间复杂度: O(n)

from typing import *

ans = [0] * 1001 # 初始化结果数组
for n in range(1, 1001): # 对每个 n 进行处理
    sq = str(n**2) # 计算 n 的平方并转为字符串

    def dfs(start, s): # 定义递归函数，尝试所有分割方式
        if start == len(sq): # 如果处理完所有字符
            return s == n # 检查当前的和是否等于 n
        res = False
        for i in range(start + 1, len(sq) + 1): # 遍历所有可能的下一个分割点
            res = dfs(i, s + int(sq[start:i])) or res # 递归检查当前分割
        return res

    ans[n] = n**2 if dfs(0, 0) else 0 # 如果能找到合适的分割，保存 n 的平方，否则为 0
for i in range(1, 1001): # 累加前缀和，方便快速查询
    ans[i] += ans[i - 1]

class Solution:
    def punishmentNumber(self, n: int) -> int: # 查询函数
        return ans[n] # 直接返回预计算结果

Explore

预计算所有可能的 n 的惩罚数而不是在每次查询时动态计算的主要原因是效率和响应速度。通过预先计算和存储结果，我们可以将查询的时间复杂度降至 O(1)，即直接通过索引访问预计算的结果，这对于需要频繁查询的情况非常有效。相反，如果每次查询时都动态计算，可能需要较复杂的递归或其他计算，特别是当 n 较大时，这将消耗更多计算资源和时间。因此，预计算为我们提供了一个时间换空间的优化策略，特别适合解答多次查询，保证查询效率。

在递归函数 `dfs` 中的终止条件 `start == len(sq)` 意味着递归函数已经遍历完了整个平方数字符串 `sq` 的所有字符，没有更多的字符可以用来继续分割。这个条件是检查是否到达字符串的末尾，如果是，则检查当前累积的子字符串的整数和是否等于原始的整数 n。确保所有子字符串被考虑的方法在于，递归函数从当前 `start` 位置向后探索所有可能的分割点，即每次调用 `dfs(i, s + int(sq[start:i]))` 都会尝试从 `start` 到 `i` 的一个新的子字符串，并将其转化为整数加到当前的和 `s` 中。这种方法通过遍历所有可能的分割点来确保覆盖所有可能的子字符串组合。

提到分割点可能性为 `2^9` 是基于字符串的最大长度和可能的分割方法计算得出的。例如，如果一个数字 n 的平方产生了一个长度为 10 的字符串（如 n=32，n^2=1024，长度为 4），那么在这个长度为 10 的字符串中，每个字符位置都可以是一个潜在的分割点。从第一个字符到第九个字符，每个位置都可以选择是否作为分割点，因此每个位置都有分割或不分割两种选择。因此，分割的可能性是 `2^9`。分割点选择的逻辑是考虑字符串中每个可能的字符位置，根据这个位置是否作为分割点，递归地探索所有不同的分割组合。