2698. Find the Punishment Number of an Integer

• categories: Code, Interview Question, leetcode, Medium
• source: https://leetcode.com/problems/find-the-punishment-number-of-an-integer
• topics: Recursion, String Manipulation

Description

Given a positive integer n, return the punishment number of n.

The punishment number of n is defined as the sum of the squares of all integers i such that:

• 1 <= i <= n
• The decimal representation of i * i can be partitioned into contiguous substrings such that the sum of the integer values of these substrings equals i.

Idea

• Use recursion

Code

class Solution:
    def punishmentNumber(self, n: int) -> int:
        def partition(prev_sum, digits, i):
            if not digits and prev_sum == i:
                return True
            for t in range(1, len(digits) + 1):
                if partition(prev_sum + int(digits[0:t]), digits[t:], i):
                    return True
            return False
        
        return sum(x for i in range(1, n+1) if partition(0, str(x:= i * i), i))