Leetcode 2698. Find the Punishment Number of an Integer

[Woodyiiiiiii]

2698. Find the Punishment Number of an Integer

2698. Find the Punishment Number of an Integer
预处理 + 回溯（Python/Java/C++/Go）

这周周赛第三题。这道题一看就知道直接暴力回溯。我还加了缓存memo，表示一个整数能对应的子串之和的集合。

class Solution {
    Map<Integer, Set<Integer>> memo = new HashMap<>();

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

    private boolean isPunishNum(String str, int val) {
        if (val < 0) {
            return false;
        }
        int all = Integer.parseInt(str);
        if (all == val) {
            return true;
        }
        if (memo.containsKey(all) && memo.get(all).contains(val)) {
            return true;
        }
        if (all < val) {
            return false;
        }
        boolean ans = false;
        for (int i = 1; i < str.length(); i++) {
            ans |= isPunishNum(str.substring(i), val - Integer.parseInt(str.substring(0, i)));
        }
        if (ans) {
            if (!memo.containsKey(all)) {
                memo.put(all, new HashSet<>());
            }
            memo.get(all).add(val);
        }
        return ans;
    }
}

分析下上面解法的时间复杂度，循环外层n，内层对于数字i，其长度是m=1+2log10(i)，所以回溯时间是m * (2^m)，最大也就是6 * (2^6)，大概100层级，所以整个循环的时间复杂度是接近100n。

可以发现，去掉memo缓存后程序运行时间没什么变化，所以这层缓存没什么用。

第二种做法是预处理。既然数组最大长度只有1000，可以预先算好所有n的答案。这样循环外层固定为1000。其次经过上面的回溯时间复杂度分析，可以不需要memo缓存。

class Solution {

    int[] preSum = new int[1001];

    public int punishmentNumber(int n) {
        for (int i = 1; i <= 1000; i++) {
            int square = i * i;
            preSum[i] = preSum[i - 1];
            String squareStr = String.valueOf(square);
            char[] squareChar = squareStr.toCharArray();
            if (isPunishNum(squareChar, i, 0, 0)) {
                preSum[i] += i * i;
            }
        }
        return preSum[n];
    }

    private boolean isPunishNum(final char[] chars, final int sum, int p, int val) {
        if (p >= chars.length) {
            return val == sum;
        }
        int x = 0;
        for (int i = p; i < chars.length; i++) {
            x = x * 10 + (chars[i] - '0');
            if (isPunishNum(chars, sum, i + 1, val + x)) {
                return true;
            }
        }
        return false;
    }
}

以上是改进后的做法，优势在于：

1. 预处理，外循环为常数
2. 取消了substring，减少了栈使用空间
3. 回溯中如果符合条件直接return true，剪枝
4. 回溯中不用parseInt，用了递加的方式计算数值