Find the punishment number of an integer

Math Backtracking

81.7% acceptance

Feb 25, 2026

1181

248

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 The decimal representation of i * i can be partitioned into contiguous substrings such that the sum of the integer values of these substrings equals i.

Solution

Rust

Time O(2^n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn punishment_number(n: i32) -> i32 {
    fn can_partition(s: &[u8], target: i32) -> bool {
      if target < 0 { return false; }
      if s.is_empty() { return target == 0; }
      let mut val = 0i32;
      for i in 0..s.len() {
        val = val * 10 + (s[i] - b'0') as i32;
        if can_partition(&s[i+1..], target - val) { return true; }
      }
      false
    }
    (1..=n).filter(|&i| {
      let sq = (i * i).to_string();
      can_partition(sq.as_bytes(), i)
    }).map(|i| i * i).sum()
  }
}