Kimi LeetCode 3405. 统计恰好有 K 个相等相邻元素的数组数目 Rust实现
题目分析
LeetCode 3405 要求统计长度为 `n`、元素范围 `[1, m]`、恰好有 `k` 对相邻元素相等的数组数目。
核心公式(组合数学):
\text{答案} = C(n-1, k) \times m \times (m-1)^{n-1-k} \pmod{10^9+7}
---
Rust 实现
方案一:预处理阶乘(推荐,O(n) 预处理)
在方法内用 `Vec` 预处理阶乘和逆阶乘,单次调用效率最高:
```rust
impl Solution {
const MOD: i64 = 1_000_000_007;
pub fn count_good_arrays(n: i32, m: i32, k: i32) -> i32 {
let n = n as usize;
let k = k as usize;
let m = m as i64;
// 快速幂
fn qpow(mut base: i64, mut exp: i64, modv: i64) -> i64 {
let mut res = 1;
base %= modv;
while exp > 0 {
if exp & 1 == 1 {
res = res * base % modv;
}
base = base * base % modv;
exp >>= 1;
}
res
}
// 预处理阶乘
let mut fact = vec![1i64; n];
for i in 1..n {
fact[i] = fact[i - 1] * i as i64 % Self::MOD;
}
// 预处理阶乘逆元(逆推)
let mut inv_fact = vec![1i64; n];
inv_fact[n - 1] = qpow(fact[n - 1], Self::MOD - 2, Self::MOD);
for i in (1..n).rev() {
inv_fact[i - 1] = inv_fact[i] * i as i64 % Self::MOD;
}
// 组合数 C(n-1, k)
let comb = fact[n - 1] * inv_fact[k] % Self::MOD * inv_fact[n - 1 - k] % Self::MOD;
let mut ans = comb * m % Self::MOD;
ans = ans * qpow(m - 1, (n - 1 - k) as i64, Self::MOD) % Self::MOD;
ans as i32
}
}
```
方案二:直接计算组合数(O(k \log \text{MOD}),无需额外空间)
如果追求代码简洁,可直接用循环计算组合数:
```rust
impl Solution {
const MOD: i64 = 1_000_000_007;
fn qpow(mut base: i64, mut exp: i64) -> i64 {
let mut res = 1;
base %= Self::MOD;
while exp > 0 {
if exp & 1 == 1 {
res = res * base % Self::MOD;
}
base = base * base % Self::MOD;
exp >>= 1;
}
res
}
fn comb(n: i32, k: i32) -> i64 {
if k < 0 || k > n {
return 0;
}
let mut res = 1i64;
for i in 0..k {
res = res * (n - i) as i64 % Self::MOD;
res = res * Self::qpow((i + 1) as i64, Self::MOD - 2) % Self::MOD;
}
res
}
pub fn count_good_arrays(n: i32, m: i32, k: i32) -> i32 {
let mut ans = Self::comb(n - 1, k);
ans = ans * m as i64 % Self::MOD;
ans = ans * Self::qpow((m - 1) as i64, (n - 1 - k) as i64) % Self::MOD;
ans as i32
}
}
```
---
两种方案对比
方案 时间复杂度 空间复杂度 特点
方案一 O(n) O(n) 预处理阶乘,查询 O(1),推荐
方案二 O(k \cdot \log \text{MOD}) O(1) 代码简洁,无额外空间
推荐方案一,虽然用了 O(n) 空间,但 Rust 的 `Vec` 管理高效,且 n \le 10^5 完全在可控范围内。
更多推荐


所有评论(0)