Count k reducible numbers less than n

Math String Dynamic Programming Combinatorics

27.6% acceptance

Feb 24, 2026

You are given a binary string s representing a number n in its binary form. You are also given an integer k. An integer x is called k-reducible if performing the following operation at most k times reduces it to 1: Replace x with the count of set bits in its binary representation. For example, the binary representation of 6 is "110". Applying the operation once reduces it to 2 (since "110" has two set bits). Applying the operation again to 2 (binary "10") reduces it to 1 (since "10" has one set bit). Return an integer denoting the number of positive integers less than n that are k-reducible. Since the answer may be too large, return it modulo 10^9 + 7.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn count_k_reducible_numbers(s: String, k: i32) -> i32 {
    const MOD: i64 = 1_000_000_007;
    let bits: Vec<u8> = s.bytes().map(|b| b - b'0').collect();
    let n = bits.len();

    // Precompute C(i, j) for i,j <= 800
    let mut comb = vec![vec![0i64; n + 1]; n + 1];
    for i in 0..=n {
      comb[i][0] = 1;
      for j in 1..=i {
        comb[i][j] = (comb[i-1][j-1] + comb[i-1][j]) % MOD;
      }
    }

    // is_k_red(c): can popcount c be reduced to 1 in at most k steps?
    // c itself is a popcount value (1..=n). We apply popcount repeatedly.
    let is_k_red = |c: usize| -> bool {
      if c == 0 { return false; }
      let mut x = c;
      let mut steps = 0;
      while x != 1 && steps < (k as usize).saturating_sub(1) {
        x = x.count_ones() as usize;
        steps += 1;
      }
      x == 1
    };

    // Digit DP: count numbers with exactly n bits (leading 1) that are < s,
    // plus all numbers with 1..(n-1) bits.
    let mut ans: i64 = 0;

    // Numbers with fewer than n bits: lengths 1..n-1
    for len in 1..n {
      // Exactly len bits: first bit is 1, remaining len-1 bits free.
      // Count with popcount = c: C(len-1, c-1)
      for c in 1..=len {
        if is_k_red(c) {
          ans = (ans + comb[len - 1][c - 1]) % MOD;
        }
      }
    }

    // Numbers with exactly n bits, < s (tight constraint from left):
    // dp[ones_so_far] under tight prefix; when we place a digit < s[i], the rest are free.
    let mut tight = vec![0i64; n + 1];
    tight[0] = 1;

    for i in 0..n {
      let mut new_tight = vec![0i64; n + 1];
      for ones_so_far in 0..=i {
        if tight[ones_so_far] == 0 { continue; }
        let b = bits[i] as usize;
        // Place digit 0..(b-1): becomes non-tight, fill remaining bits freely
        for d in 0..b {
          // at i=0, d=0 would mean a number with fewer bits -> skip (already counted above)
          if i == 0 && d == 0 { continue; }
          let new_ones = ones_so_far + d;
          let rem = n - 1 - i;
          // choose any extra ones in remaining rem bits
          for extra in 0..=rem {
            let total = new_ones + extra;
            if total >= 1 && is_k_red(total) {
              ans = (ans + tight[ones_so_far] * comb[rem][extra]) % MOD;
            }
          }
        }
        // Place digit b: stays tight
        new_tight[ones_so_far + b] = (new_tight[ones_so_far + b] + tight[ones_so_far]) % MOD;
      }
      tight = new_tight;
    }
    // tight path represents n itself — excluded since we want strictly less than n

    ans as i32
  }
}