题目分析

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 完全在可控范围内。

 

更多推荐