<p>Given a positive integer <code>n</code>, return <em>the <strong>punishment number</strong></em> of <code>n</code>.</p>

<p>The <strong>punishment number</strong> of <code>n</code> is defined as the sum of the squares of all integers <code>i</code> such that:</p>

<ul>
	<li><code>1 &lt;= i &lt;= n</code></li>
	<li>The decimal representation of <code>i * i</code> can be partitioned into contiguous substrings such that the sum of the integer values of these substrings equals <code>i</code>.</li>
</ul>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> n = 10
<strong>Output:</strong> 182
<strong>Explanation:</strong> There are exactly 3 integers i 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.
Hence, the punishment number of 10 is 1 + 81 + 100 = 182
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> n = 37
<strong>Output:</strong> 1478
<strong>Explanation:</strong> There are exactly 4 integers i 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
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= n &lt;= 1000</code></li>
</ul>

import math
class Solution:
    def punishmentNumber(self, n: int) -> int:

        # memory = {}
        total_res = 0

        def partitions(nb):

            """
            1  -> 1
            12 -> [12], [1,2]
            123 -> [123], [1, 23], [12, 3], [1, 2, 3] || [12], [1,2] + [3]-> [123], [12, 3], [1,23], [1,2,3] 
            1234 -> || [123], [12, 3], [1,23], [1,2,3] + [4] -> [1234], [123,4], [12, 34], [123,4], [1,234], [1,23,4], [1,2,34], [1,2,3,4]
            """
            
            # if nb in memory:
            #     return memory[nb]

            if nb < 10:
                return [[nb]]
            else:
                current_digit = nb % 10
                nb = nb//(10)
                res = []
                for part in partitions(nb):
                    res.append(part + [current_digit])
                    res.append(part[:-1] + [current_digit + part[-1]*10])
                # memory[nb] = res
                return res

      
        for i in range(1, n+1):
            res = partitions(i**2)
            for tmp_res in res:
                if sum(tmp_res) == i:
                    total_res += i**2
                    break
        return total_res
    


            
            

376