2025-07-14:统计恰好有 K 个相等相邻元素的数组数目。用go语言,给定三个整数 n、m、k,定义一个长度为 n 的数组 arr 满足以下条件:
o arr 中的每个元素都是 1 到 m 之间的整数(包含边界)。
o 在数组中恰好存在 k 个位置 i(1 <= i < n),使得 arr[i - 1] 和 arr[i] 相等。
请计算满足上述条件的不同数组 arr 的数量。
由于结果可能非常大,请将最终答案对 1000000007 取模后返回。
1 <= n <= 100000。
1 <= m <= 100000。
0 <= k <= n - 1。
输入:n = 3, m = 2, k = 1。
输出:4。
解释:
总共有 4 个好数组,分别是 [1, 1, 2] ,[1, 2, 2] ,[2, 1, 1] 和 [2, 2, 1] 。
所以答案为 4 。
题目来自力扣3405。
解决思路
1. 问题分解:
o 我们需要构造一个长度为 n 的数组,其中有恰好 k 个相邻元素对是相等的。
o 这可以转化为:将数组分成 n - k 个“块”,其中每个块内的元素相同,但相邻块之间的元素不同。
o 例如,n = 3,k = 1 时,数组可以分成 2 个块(因为 n - k = 2),如 [1, 1, 2](块 [1, 1] 和 [2])。
2. 组合数学:
o 首先,我们需要选择 k 个位置作为相邻相等的位置。这相当于在 n - 1 个可能的位置中选择 k 个位置,使得这些位置是相邻相等的。这可以通过组合数 C(n - 1, k) 来计算。
o 然后,我们需要为这些“块”分配值。第一个块有 m 种选择(1 到 m),后续的每个新块必须与前一个块的值不同,因此有 m - 1 种选择。
o 总共有 n - k 个块(因为 n 个元素分成 n - k 个块),因此选择的方案数为 m * (m - 1)^(n - k - 1)。
3. 组合公式:
o 最终的数量为 C(n - 1, k) * m * (m - 1)^(n - k - 1)。
o 例如,n = 3,m = 2,k = 1:
o C(2, 1) = 2。
o m = 2。
o (m - 1)^(n - k - 1) = 1^1 = 1。
o 总数为 2 * 2 * 1 = 4,与示例一致。
具体步骤
1. 预计算阶乘和逆阶乘:
o 为了高效计算组合数 C(n, k),我们需要预计算阶乘 fact[i] = i! mod MOD 和逆阶乘 invFact[i] = (i!)^(-1) mod MOD。
o 阶乘可以通过递推计算:fact[i] = fact[i - 1] * i % MOD。
o 逆阶乘可以通过费马小定理计算:invFact[i] = fact[i]^(MOD - 2) mod MOD(因为 MOD 是质数)。
o 逆阶乘的递推计算:invFact[i - 1] = invFact[i] * i % MOD。
2. 计算组合数:
o C(n, k) = fact[n] * invFact[k] % MOD * invFact[n - k] % MOD。
3. 计算结果:
o 使用预计算的组合数和快速幂计算结果:
o ans = C(n - 1, k) * m % MOD * (m - 1)^(n - k - 1) % MOD。
时间和空间复杂度
1. 时间复杂度:
o 预计算阶乘和逆阶乘的时间是 O(MX),其中 MX 是预计算的最大值(这里是 100000)。
o 计算组合数和快速幂的时间是 O(1)(因为预计算已经完成)。
o 因此,总的时间复杂度是 O(MX),即 O(100000)。
2. 空间复杂度:
o 需要存储阶乘和逆阶乘数组 fact 和 invFact,大小均为 MX。
o 因此,额外的空间复杂度是 O(MX),即 O(100000)。
总结
o 通过组合数学将问题分解为选择相邻相等的位置和分配块的值。
o 预计算阶乘和逆阶乘以高效计算组合数。
o 最终结果是组合数与块值分配方案的乘积。
o 时间和空间复杂度均为 O(MX),其中 MX 是预计算的最大值(100000)。
Go完整代码如下:
.
package main
import (
"fmt"
)
const MOD = 1e9 + 7
const MX = 100000
var fact = make([]int64, MX)
var invFact = make([]int64, MX)
func qpow(x int64, n int)int64 {
res := int64(1)
for n > 0 {
if n&1 == 1 {
res = res * x % MOD
}
x = x * x % MOD
n >>= 1
}
return res
}
func initFact() {
if fact[0] != 0 {
return
}
fact[0] = 1
for i := 1; i < MX; i++ {
fact[i] = fact[i-1] * int64(i) % MOD
}
invFact[MX-1] = qpow(fact[MX-1], MOD-2)
for i := MX - 1; i > 0; i-- {
invFact[i-1] = invFact[i] * int64(i) % MOD
}
}
func comb(n, m int)int64 {
return fact[n] * invFact[m] % MOD * invFact[n-m] % MOD
}
func countGoodArrays(n, m, k int)int {
initFact()
returnint(comb(n-1, k) * int64(m) % MOD * qpow(int64(m-1), n-k-1) % MOD)
}
func main() {
n := 3
m := 2
k := 1
result := countGoodArrays(n, m, k)
fmt.Println(result)
}
Python完整代码如下:
.
# -*-coding:utf-8-*-
MOD = 10**9 + 7
MX = 100000
fact = [0] * MX
invFact = [0] * MX
def qpow(x: int, n: int) -> int:
res = 1
base = x % MOD
while n > 0:
if n & 1:
res = res * base % MOD
base = base * base % MOD
n >>= 1
return res
def initFact():
if fact[0] != 0:
return
fact[0] = 1
for i in range(1, MX):
fact[i] = fact[i-1] * i % MOD
invFact[MX-1] = qpow(fact[MX-1], MOD-2)
for i in range(MX-2, -1, -1):
invFact[i] = invFact[i+1] * (i+1) % MOD
def comb(n: int, m: int) -> int:
if m > n or m < 0:
return0
return fact[n] * invFact[m] % MOD * invFact[n-m] % MOD
def countGoodArrays(n: int, m: int, k: int) -> int:
initFact()
return comb(n-1, k) * m % MOD * qpow(m-1, n-k-1) % MOD
if __name__ == '__main__':
n = 3
m = 2
k = 1
print(countGoodArrays(n, m, k))·
我们相信Go语言和算法为普通开发者提供了困境的“面试利器”,并致力于分享全面的编程知识。在这里,您可以找到最新的Go语言教程、算法解析、提升面试竞争力的秘籍以及行业动态。
欢迎关注“福大规模架构师每日一题”,让 Go 语言和算法助力您的职业发展
·
