diff options
author | Jack Sangdahl <[email protected]> | 2025-02-14 22:01:10 -0700 |
---|---|---|
committer | Jack Sangdahl <[email protected]> | 2025-02-14 22:01:10 -0700 |
commit | d2b98c72a470837443f9e0e5d09a2c924c760353 (patch) | |
tree | 3315e71f989e1700cb870074105578f8726c7564 /find_punishment_number_integer.py | |
parent | 174c3b55b04c6a71e22712ce599afdece018ee39 (diff) | |
download | leetcode-d2b98c72a470837443f9e0e5d09a2c924c760353.tar.gz leetcode-d2b98c72a470837443f9e0e5d09a2c924c760353.zip |
2698. find the punishment number of an integer
Diffstat (limited to 'find_punishment_number_integer.py')
-rw-r--r-- | find_punishment_number_integer.py | 42 |
1 files changed, 42 insertions, 0 deletions
diff --git a/find_punishment_number_integer.py b/find_punishment_number_integer.py new file mode 100644 index 0000000..cc66059 --- /dev/null +++ b/find_punishment_number_integer.py @@ -0,0 +1,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)) |