Number of integers with popcount depth equal to k ii

Array Divide and Conquer Binary Indexed Tree Segment Tree

59.4% acceptance

Feb 25, 2026

You are given an integer array nums. For any positive integer x, define the following sequence: p0 = x pi+1 = popcount(pi) for all i >= 0, where popcount(y) is the number of set bits (1's) in the binary representation of y. This sequence will eventually reach the value 1. The popcount-depth of x is defined as the smallest integer d >= 0 such that pd = 1. You are also given a 2D integer array queries, where each queries[i] is either: [1, l, r, k] - Determine the number of indices j such that l <= j <= r and the popcount-depth of nums[j] is equal to k. [2, idx, val] - Update nums[idx] to val. Return an integer array answer, where answer[i] is the number of indices for the ith query of type [1, l, r, k].

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn popcount_depth(nums: Vec<i64>, queries: Vec<Vec<i64>>) -> Vec<i32> {
    fn depth(mut x: i64) -> usize {
      let mut d = 0usize;
      while x != 1 {
        x = x.count_ones() as i64;
        d += 1;
      }
      d
    }

    let n = nums.len();
    let mut depths: Vec<usize> = nums.iter().map(|&v| depth(v)).collect();

    // BIT (Fenwick tree) for each depth 0..=5
    const MAXK: usize = 6;
    let mut bits = vec![vec![0i32; n + 1]; MAXK];

    fn update(bit: &mut Vec<i32>, mut i: usize, delta: i32) {
      i += 1;
      while i < bit.len() {
        bit[i] += delta;
        i += i & i.wrapping_neg();
      }
    }
    fn query(bit: &Vec<i32>, mut i: usize) -> i32 {
      let mut s = 0;
      i += 1;
      while i > 0 {
        s += bit[i];
        i -= i & i.wrapping_neg();
      }
      s
    }

    for i in 0..n {
      let d = depths[i];
      if d < MAXK { update(&mut bits[d], i, 1); }
    }

    let mut result = Vec::new();
    for q in &queries {
      if q[0] == 1 {
        let (l, r, k) = (q[1] as usize, q[2] as usize, q[3] as usize);
        if k < MAXK {
          let cnt = query(&bits[k], r) - if l > 0 { query(&bits[k], l - 1) } else { 0 };
          result.push(cnt);
        } else {
          result.push(0);
        }
      } else {
        let (idx, val) = (q[1] as usize, q[2]);
        let old_d = depths[idx];
        let new_d = depth(val);
        if old_d < MAXK { update(&mut bits[old_d], idx, -1); }
        if new_d < MAXK { update(&mut bits[new_d], idx, 1); }
        depths[idx] = new_d;
      }
    }
    result
  }
}