# Problem Statement 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`. ## Constraints - `1 <= n <= 1000` # Solution This problem can be solved by computing all punishment numbers and adding them to a [[set]]. We then can check all the numbers up to `n` to see if they are in the set, and if so, add their square to some accumulator variable. ```python return sum(i**2 for i in range(1, n+1) if i in good) ``` Where it gets tricky is in the punishment number computation. There are a few ways to do this, like DP or backtracking, but I opted to use a [[Bitmask|bitmask]]. Essentially, we iterate through all subsets of the digits of `i*i`, which will give us an array of `True` and `False` , representing whether that digit in included in that subset. However, we can also use this array to generate all partitionings by considering the continuous groups. Using this boolean array, I [[zip]] it with the string representation of the number and use [[groupby]] to extract the groups. Once we have the continuous groups, I add the integer representation of the [[Python join|joined]] digits for each group to some variable and then check if the total sum is equal to `i*i`. If so, I add it to the `good` set. ```python good = set() for i in range(1, 1001): s = str(i*i)     k = len(s)     for mask in range(2**k):         total = 0         conv = [bool(mask & (2**j)) for j in range(k)]         g = groupby(zip(s, conv), key=lambda x: x[1])         for key, val in g:             total += int("".join([v[0] for v in val]))         if total == i:             good.add(i) ``` Time Complexity: O$(n*2^{\log_{}n})$ | Space Complexity: O$(n)$