blob: cc66059af05e43d1202da48252baa8ed67b6dcdb (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
|
# 2698. Find the Punishment Number of an Integer
from bisect import bisect_right
from itertools import accumulate
"""
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' and the decimal representation of 'i * i' can be
partitioned into contiguous substrings such that the sum of the integer
values of the substrings equals 'i'.
"""
sqr = lambda x: x * x
f = lambda x: x % 9 < 2 and valid(x)
nums = tuple(filter(f, tuple(range(1, 1001))))
punish = tuple(accumulate(map(sqr, nums), initial=0))
class Solution:
def valid(self, n):
sq = str(sqr(n))
for i in range(1 << (len(sq) - 1)):
s = [str()]
for j, ch in enumerate(sq):
s[-1] += ch
if i & (1 << j):
s.append(str())
if sum(map(int, s)) == n:
return True
return False
def punishmentNumber(self, n):
return punish[bisect_right(nums, n)]
if __name__ == "__main__":
obj = Solution()
print(obj.punishmentNumber(10))
print(obj.punishmentNumber(37))
|