Find the Punishment Number of an Integer

lcpy gen -n 2698   # by problem number
lcpy gen -s find_the_punishment_number_of_an_integer   # by problem name

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

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

class Solution:
    # Time: O(n * d^2) where d is the digit count of i*i (partition search per i)
    # Space: O(d) recursion depth
    def punishment_number(self, n: int) -> int:
        def can_partition(sq: str, target: int, idx: int = 0, cur: int = 0) -> bool:
            if idx == len(sq):
                return cur == target
            for j in range(idx + 1, len(sq) + 1):
                part = int(sq[idx:j])
                if cur + part > target:
                    break
                if can_partition(sq, target, j, cur + part):
                    return True
            return False

        total = 0
        for i in range(1, n + 1):
            if can_partition(str(i * i), i):
                total += i * i
        return total