Kimi LeetCode 3621. 位计数深度为 K 的整数数目 I Rust实现

LeetCode 3621. 位计数深度为 K 的整数数目 I Rust 实现

```rust
impl Solution {
pub fn popcount_depth(n: i64, k: i32) -> i64 {
// k=0：只有 1 的深度为 0
if k == 0 {
return if n >= 1 { 1 } else { 0 };
}

// 计算一个整数的 popcount
fn pc(mut x: i32) -> i32 {
let mut cnt = 0;
while x > 0 {
cnt += x & 1;
x >>= 1;
}
cnt
}

// 预处理 1~60 的 popcount-depth
let mut depth = vec![-1; 65];
depth[1] = 0;

fn get_depth(x: i32, depth: &mut Vec<i32>) -> i32 {
if x == 1 {
return 0;
}
if depth[x as usize] != -1 {
return depth[x as usize];
}
let d = 1 + get_depth(pc(x), depth);
depth[x as usize] = d;
d
}

// 找出所有满足 depth(c) == k-1 的 popcount 值 c
let mut targets = vec![];
for i in 1..=60 {
if get_depth(i, &mut depth) == k - 1 {
targets.push(i);
}
}

if targets.is_empty() {
return 0;
}

// n 的二进制表示，按位存储
let bits: Vec<u8> = format!("{:b}", n).bytes().map(|b| b - b'0').collect();
let m = bits.len();
let mut ans = 0i64;

for &c in &targets {
// memo[pos][ones]，仅非 tight 状态有效，-1 表示未计算
let mut memo = vec![vec![-1i64; (c + 1) as usize]; m];

fn dfs(
pos: usize,
ones: i32,
tight: bool,
bits: &Vec<u8>,
c: i32,
memo: &mut Vec<Vec<i64>>,
) -> i64 {
if ones > c {
return 0; // 1 的个数已超，剪枝
}
if pos == bits.len() {
return if ones == c { 1 } else { 0 };
}

if !tight && memo[pos][ones as usize] != -1 {
return memo[pos][ones as usize];
}

let limit = if tight { bits[pos] } else { 1 };
let mut res = 0i64;

for d in 0..=limit {
res += dfs(
pos + 1,
ones + d as i32,
tight && (d == limit),
bits,
c,
memo,
);
}

if !tight {
memo[pos][ones as usize] = res;
}
res
}

ans += dfs(0, 0, true, &bits, c, &mut memo);
}

// k=1 的特殊处理：popcount=1 包含 x=1，但 depth(1)=0，需排除
if k == 1 {
ans -= 1;
}

ans
}
}
```

---

Rust 实现要点

要点 说明
`format!("{:b}", n)` 将 `n` 转为二进制字符串，再拆分为 `Vec<u8>` 逐位处理
内部递归函数 `dfs` 作为独立函数接收所有必要参数（闭包递归在 Rust 中较繁琐）
记忆化数组 `memo[pos][ones]` 用 `-1` 标记未计算，仅 `tight = false` 时写入复用
`i64` 返回值 `n ≤ 10¹⁵`，答案在 `i64` 范围内；`pc` 用 `i32` 即可（`x ≤ 60`）
k=1 修正 `c = 1` 时会把 `x = 1` 算入，但其深度为 0，需减去

时间复杂度： `O(log² n)`

空间复杂度： `O(log² n)`