aboutsummaryrefslogtreecommitdiff
path: root/find_punishment_number_integer.py
diff options
context:
space:
mode:
authorJack Sangdahl <[email protected]>2025-02-14 22:01:10 -0700
committerJack Sangdahl <[email protected]>2025-02-14 22:01:10 -0700
commitd2b98c72a470837443f9e0e5d09a2c924c760353 (patch)
tree3315e71f989e1700cb870074105578f8726c7564 /find_punishment_number_integer.py
parent174c3b55b04c6a71e22712ce599afdece018ee39 (diff)
downloadleetcode-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.py42
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))