diff --git a/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/README.md b/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/README.md index 3a69d275be3eb..26286cd16e64c 100644 --- a/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/README.md +++ b/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/README.md @@ -94,7 +94,23 @@ tags: -### 方法一 +### 方法一：组合数学 + 快速幂 + +长度为 $n$ 的数组，一共有 $n - 1$ 个相邻元素对。我们需要从这 $n - 1$ 个相邻元素对中选出 $k$ 个，使得这 $k$ 个相邻元素对的两个元素相等，那么剩下的 $n - 1 - k$ 个相邻元素对的两个元素不相等。 + +这相当于我们对数组进行 $n - 1 - k$ 次分割，得到 $n - k$ 个分段，每个分段子数组中的元素都相等，分割方案为 $C_{n - 1}^{n - 1 - k} = C_{n - 1}^{k}$。 + +第一段，我们可以任意选择一个元素，即在 $[1, m]$ 中选择一个元素，剩下的 $n - k - 1$ 个分段，只要确保每个分段的元素都不等于前一个分段的元素即可，因此剩下的每个分段都有 $m - 1$ 种选择，一共有 $m \times (m - 1)^{n - k - 1}$ 种选择。 + +将上述两部分结合起来，我们可以得到答案为： + +$$ +C_{n - 1}^{k} \times m \times (m - 1)^{n - k - 1} \bmod (10^9 + 7) +$$ + +在代码实现上，我们可以预处理阶乘和逆元，使用快速幂计算组合数。 + +忽略预处理的时间和空间，时间复杂度 $O(\log (n - k))$，空间复杂度 $O(1)$。 @@ -278,6 +294,67 @@ export function countGoodArrays(n: number, m: number, k: number): number { } ``` +#### Rust + +```rust +impl Solution { + pub fn count_good_arrays(n: i32, m: i32, k: i32) -> i32 { + const N: usize = 1e5 as usize + 10; + const MOD: i64 = 1_000_000_007; + use std::sync::OnceLock; + + static F: OnceLock> = OnceLock::new(); + static G: OnceLock> = OnceLock::new(); + + fn qpow(mut a: i64, mut k: i64, m: i64) -> i64 { + let mut res = 1; + while k != 0 { + if k & 1 == 1 { + res = res * a % m; + } + a = a * a % m; + k >>= 1; + } + res + } + + fn init() -> (&'static Vec, &'static Vec) { + F.get_or_init(|| { + let mut f = vec![1i64; N]; + for i in 1..N { + f[i] = f[i - 1] * i as i64 % MOD; + } + f + }); + + G.get_or_init(|| { + let f = F.get().unwrap(); + let mut g = vec![1i64; N]; + for i in 1..N { + g[i] = qpow(f[i], MOD - 2, MOD); + } + g + }); + + (F.get().unwrap(), G.get().unwrap()) + } + + fn comb(f: &[i64], g: &[i64], m: usize, n: usize) -> i64 { + f[m] * g[n] % MOD * g[m - n] % MOD + } + + let (f, g) = init(); + let n = n as usize; + let m = m as i64; + let k = k as usize; + + let c = comb(f, g, n - 1, k); + let pow = qpow(m - 1, (n - 1 - k) as i64, MOD); + (c * m % MOD * pow % MOD) as i32 + } +} +``` + diff --git a/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/README_EN.md b/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/README_EN.md index 520176390ed1d..0160b9c3ab4e6 100644 --- a/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/README_EN.md +++ b/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/README_EN.md @@ -90,7 +90,23 @@ tags: -### Solution 1 +### Solution 1: Combinatorics + Fast Power + +For an array of length $n$, there are $n - 1$ pairs of adjacent elements. We need to select $k$ of these $n - 1$ adjacent pairs such that the two elements in each of these $k$ pairs are equal, and the remaining $n - 1 - k$ adjacent pairs have different elements. + +This is equivalent to splitting the array $n - 1 - k$ times, resulting in $n - k$ segments, where all elements in each segment are equal. The number of ways to split is $C_{n - 1}^{n - 1 - k} = C_{n - 1}^{k}$. + +For the first segment, we can choose any element from $[1, m]$. For the remaining $n - k - 1$ segments, we just need to ensure that the element in each segment is different from the previous segment, so each of these segments has $m - 1$ choices. In total, there are $m \times (m - 1)^{n - k - 1}$ ways to choose. + +Combining the two parts above, we get the answer: + +$$ +C_{n - 1}^{k} \times m \times (m - 1)^{n - k - 1} \bmod (10^9 + 7) +$$ + +In the code implementation, we can precompute factorials and inverses, and use fast power to calculate combinations. + +Ignoring the preprocessing time and space, the time complexity is $O(\log (n - k))$, and the space complexity is $O(1)$. @@ -274,6 +290,67 @@ export function countGoodArrays(n: number, m: number, k: number): number { } ``` +#### Rust + +```rust +impl Solution { + pub fn count_good_arrays(n: i32, m: i32, k: i32) -> i32 { + const N: usize = 1e5 as usize + 10; + const MOD: i64 = 1_000_000_007; + use std::sync::OnceLock; + + static F: OnceLock> = OnceLock::new(); + static G: OnceLock> = OnceLock::new(); + + fn qpow(mut a: i64, mut k: i64, m: i64) -> i64 { + let mut res = 1; + while k != 0 { + if k & 1 == 1 { + res = res * a % m; + } + a = a * a % m; + k >>= 1; + } + res + } + + fn init() -> (&'static Vec, &'static Vec) { + F.get_or_init(|| { + let mut f = vec![1i64; N]; + for i in 1..N { + f[i] = f[i - 1] * i as i64 % MOD; + } + f + }); + + G.get_or_init(|| { + let f = F.get().unwrap(); + let mut g = vec![1i64; N]; + for i in 1..N { + g[i] = qpow(f[i], MOD - 2, MOD); + } + g + }); + + (F.get().unwrap(), G.get().unwrap()) + } + + fn comb(f: &[i64], g: &[i64], m: usize, n: usize) -> i64 { + f[m] * g[n] % MOD * g[m - n] % MOD + } + + let (f, g) = init(); + let n = n as usize; + let m = m as i64; + let k = k as usize; + + let c = comb(f, g, n - 1, k); + let pow = qpow(m - 1, (n - 1 - k) as i64, MOD); + (c * m % MOD * pow % MOD) as i32 + } +} +``` + diff --git a/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/Solution.rs b/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/Solution.rs new file mode 100644 index 0000000000000..d962f0b4eaa75 --- /dev/null +++ b/solution/3400-3499/3405.Count the Number of Arrays with K Matching Adjacent Elements/Solution.rs @@ -0,0 +1,56 @@ +impl Solution { + pub fn count_good_arrays(n: i32, m: i32, k: i32) -> i32 { + const N: usize = 1e5 as usize + 10; + const MOD: i64 = 1_000_000_007; + use std::sync::OnceLock; + + static F: OnceLock> = OnceLock::new(); + static G: OnceLock> = OnceLock::new(); + + fn qpow(mut a: i64, mut k: i64, m: i64) -> i64 { + let mut res = 1; + while k != 0 { + if k & 1 == 1 { + res = res * a % m; + } + a = a * a % m; + k >>= 1; + } + res + } + + fn init() -> (&'static Vec, &'static Vec) { + F.get_or_init(|| { + let mut f = vec![1i64; N]; + for i in 1..N { + f[i] = f[i - 1] * i as i64 % MOD; + } + f + }); + + G.get_or_init(|| { + let f = F.get().unwrap(); + let mut g = vec![1i64; N]; + for i in 1..N { + g[i] = qpow(f[i], MOD - 2, MOD); + } + g + }); + + (F.get().unwrap(), G.get().unwrap()) + } + + fn comb(f: &[i64], g: &[i64], m: usize, n: usize) -> i64 { + f[m] * g[n] % MOD * g[m - n] % MOD + } + + let (f, g) = init(); + let n = n as usize; + let m = m as i64; + let k = k as usize; + + let c = comb(f, g, n - 1, k); + let pow = qpow(m - 1, (n - 1 - k) as i64, MOD); + (c * m % MOD * pow % MOD) as i32 + } +}