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 | |
parent | 174c3b55b04c6a71e22712ce599afdece018ee39 (diff) | |
download | leetcode-d2b98c72a470837443f9e0e5d09a2c924c760353.tar.gz leetcode-d2b98c72a470837443f9e0e5d09a2c924c760353.zip |
2698. find the punishment number of an integer
-rw-r--r-- | find_punishment_number_integer.c | 62 | ||||
-rw-r--r-- | find_punishment_number_integer.py | 42 |
2 files changed, 104 insertions, 0 deletions
diff --git a/find_punishment_number_integer.c b/find_punishment_number_integer.c new file mode 100644 index 0000000..9f057fe --- /dev/null +++ b/find_punishment_number_integer.c @@ -0,0 +1,62 @@ +// 2698. Find the Punishment Number of an Integer +#include "leetcode.h" + +/* + * 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'. + */ + +bool partition(char *s, int idx, int target) { + int sum = 0, base = 1; + bool res = false; + if (!s[idx]) + return !target; + for (int i = idx; s[i]; i++) { + sum += (s[i] - '0') * base; + res = partition(s, i + 1, target - sum); + if (res) + break; + base *= 10; + } + return res; +} + +void reverse(char *s, int i, int j) { + for (; i < j; i++, j--) { + char tmp = s[i]; + s[i] = s[j]; + s[j] = tmp; + } +} + +int punishment(int n, int *dp) { + int sum = 0, res = 0; + char s[8]; + if (!n) + return 0; + if (dp[n] != -1) + return dp[n]; + sprintf(s, "%d", n * n); + reverse(s, 0, strlen(s) - 1); + res = partition(s, 0, n) ? (n * n) + punishment(n - 1, dp) + : punishment(n - 1, dp); + dp[n] = res; + return res; +} + +int punishmentNumber(int n) { + int dp[1001] = {0}; + if (!dp[0]) { + memset(dp, -1, sizeof(dp)); + punishment(1000, dp); + } + return dp[n]; +} + +int main() { + printf("%d\n", punishmentNumber(10)); // expect: 182 + printf("%d\n", punishmentNumber(37)); // expect: 1478 +} 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)) |