LeetCode/kt/find-the-punishment-number-of-an-integer.kt

class Solution {
    companion object {
        private val PARTITIONS: List<Int> = listOf(10, 100, 1000)
    }

    private fun canPartition(
        num: Int,
        target: Int,
    ): Boolean =
        when {
            target < 0 || num < target -> false
            num == target -> true
            else ->
                PARTITIONS.any { it ->
                    canPartition(num / it, target - num % it)
                }
        }

    fun punishmentNumber(n: Int): Int =
        (1..n).asSequence().sumOf {
            val square = it * it

            when {
                canPartition(square, it) -> square
                else -> 0
            }
        }
}