LeetCode in Kotlin

2698. Find the Punishment Number of an Integer

Medium

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:

Example 1:

Input: n = 10

Output: 182

Explanation: There are exactly 3 integers i that satisfy the conditions in the statement:

Hence, the punishment number of 10 is 1 + 81 + 100 = 182

Example 2:

Input: n = 37

Output: 1478

Explanation: There are exactly 4 integers i that satisfy the conditions in the statement:

Hence, the punishment number of 37 is 1 + 81 + 100 + 1296 = 1478

Constraints:

Solution

class Solution {
    fun punishmentNumber(n: Int): Int {
        fun partition(x: Int, target: Int): Boolean {
            if (x == target) return true
            if (target < 0 || x < target) return false
            return partition(x / 10, target - (x % 10)) ||
                partition(x / 100, target - (x % 100)) ||
                partition(x / 1000, target - (x % 1000))
        }
        var res = 0
        for (i in 1..n) {
            val iSq = i * i
            if (partition(iSq, i)) {
                res += iSq
            }
        }
        return res
    }
}