Today’s problem

https://leetcode.com/problems/find-the-punishment-number-of-an-integer/description/

Intuition

Note that all the combinations of a number that is less than $1000^2$ is $2^5 = 32$. Which means that if we use a dfs to decide whether to choose to break from one interval or not cost at most 32 for one number.

Since we only have $N \leq 1000$, we can solve it with brute method.

Approach

Use a dfs to check whether it cawn be a punishment number or not.

Complexity

  • Time complexity: $O(N\times 2^{log(N)})$, N is the same representation as the description.

  • Space complexity: $O(1)$.

Code

class Solution:
    def punishmentNumber(self, n: int) -> int:
        def check(x, now, s, nows, cnt):
            if now == 0:
                return (s + nows) == x
            if s > x:
                return False
            flag = check(x, now // 10, s, nows + now % 10 * (10 ** cnt), cnt + 1)
            if flag:
                return True
            flag |= check(x, now // 10, s + nows, now % 10, 1) 

            return flag

        ans = 0
        for i in range(n + 1):
            if check(i, i*i, 0, 0, 0):
                ans += i * i

        return ans

class Solution:
    def punishmentNumber(self, n: int) -> int:
        punishmentNumbers = [0, 1, 9, 10, 36, 45, 55, 82, 91, 99, 100, 235, 297, 369, 370, 379, 414, 657, 675, 703, 756, 792, 909, 918, 945, 964, 990, 991, 999, 1000]

        ans = 0
        for x in punishmentNumbers:
            if x > n:
                return ans
            ans += x * x
        
        return ans

Result of normal solution:

image.png

Result of fastest solution:

image.png