Find the Punishment Number of an Integer

This approach iterates through each integer i from 1 to n. For each i, it first calculates its square, i*i, and converts it into a string. Then, a recursive backtracking function is used to determine if this string can be partitioned into substrings whose integer values sum up to i. If a valid partition exists, i*i is added to a running total. This method directly translates the problem statement into code.

• Initialize punishmentSum = 0.
• For i from 1 to n:
  • Let s = String.valueOf(i * i).
  • If a helper function canPartition(s, i) returns true, add i * i to punishmentSum.
• Return punishmentSum.
• The canPartition function uses a recursive depth-first search (dfs) to check the condition.
• dfs(s, target, index, currentSum) function:
  • Base Case: If index reaches the end of the string s, it means we have processed the entire string. Return true if currentSum equals target, otherwise false.
  • Recursive Step: Iterate through all possible next substrings starting from index. For each substring:
    • Parse it to an integer num.
    • If currentSum + num exceeds target, we can prune this path and all subsequent paths from this index because numbers will only get larger.
    • Make a recursive call dfs(s, target, j + 1, currentSum + num) for the rest of the string.
    • If any recursive call returns true, it means a valid partition is found, so propagate true up the call stack.
  • If the loop completes without finding a valid partition, return false.

The main function punishmentNumber(n) initializes a sum to zero and then loops from i = 1 to n. Inside the loop, it calculates square = i * i and converts it to a String s. It then calls a helper function, canPartition(s, i), to check if s can be partitioned to sum to i. If canPartition returns true, square is added to the sum.

The canPartition function itself is a wrapper that initiates a recursive search. It calls a backtracking function dfs(s, target, index, currentSum) which explores all possible partitions of the string s starting from a given index to see if they can sum up to the target.

class Solution {
    public int punishmentNumber(int n) {
        int totalPunishment = 0;
        for (int i = 1; i <= n; i++) {
            String s = String.valueOf(i * i);
            if (canPartition(s, i)) {
                totalPunishment += i * i;
            }
        }
        return totalPunishment;
    }

    private boolean canPartition(String s, int target) {
        return dfs(s, target, 0, 0);
    }

    private boolean dfs(String s, int target, int index, int currentSum) {
        if (index == s.length()) {
            return currentSum == target;
        }

        for (int j = index; j < s.length(); j++) {
            String sub = s.substring(index, j + 1);
            int num = Integer.parseInt(sub);

            if (currentSum + num > target) {
                // Optimization: if current sum already exceeds target,
                // no need to check longer numbers from this index.
                break;
            }

            if (dfs(s, target, j + 1, currentSum + num)) {
                return true;
            }
        }

        return false;
    }
}