2025-07-14：统计恰好有 K 个相等相邻元素的数组数目。用go语言，给定三个整数 n、m、k，定义一个长度为 n 的数组 arr 满足以下条件：

• arr 中的每个元素都是 1 到 m 之间的整数（包含边界）。

• 在数组中恰好存在 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. 问题分解：

   • 我们需要构造一个长度为 n 的数组，其中有恰好 k 个相邻元素对是相等的。
   • 这可以转化为：将数组分成 n - k 个“块”，其中每个块内的元素相同，但相邻块之间的元素不同。
   • 例如，n = 3，k = 1 时，数组可以分成 2 个块（因为 n - k = 2），如 [1, 1, 2]（块 [1, 1] 和 [2]）。

2. 组合数学：

   • 首先，我们需要选择 k 个位置作为相邻相等的位置。这相当于在 n - 1 个可能的位置中选择 k 个位置，使得这些位置是相邻相等的。这可以通过组合数 C(n - 1, k) 来计算。
   • 然后，我们需要为这些“块”分配值。第一个块有 m 种选择（1 到 m），后续的每个新块必须与前一个块的值不同，因此有 m - 1 种选择。
   • 总共有 n - k 个块（因为 n 个元素分成 n - k 个块），因此选择的方案数为 m * (m - 1)^(n - k - 1)。

3. 组合公式：

   • 最终的数量为 C(n - 1, k) * m * (m - 1)^(n - k - 1)。
   • 例如，n = 3，m = 2，k = 1：
     • C(2, 1) = 2。
     • m = 2。
     • (m - 1)^(n - k - 1) = 1^1 = 1。
     • 总数为 2 * 2 * 1 = 4，与示例一致。

具体步骤

1. 预计算阶乘和逆阶乘：

   • 为了高效计算组合数 C(n, k)，我们需要预计算阶乘 fact[i] = i! mod MOD 和逆阶乘 invFact[i] = (i!)^(-1) mod MOD。
   • 阶乘可以通过递推计算：fact[i] = fact[i - 1] * i % MOD。
   • 逆阶乘可以通过费马小定理计算：invFact[i] = fact[i]^(MOD - 2) mod MOD（因为 MOD 是质数）。
   • 逆阶乘的递推计算：invFact[i - 1] = invFact[i] * i % MOD。

2. 计算组合数：

   • C(n, k) = fact[n] * invFact[k] % MOD * invFact[n - k] % MOD。

3. 计算结果：

   • 使用预计算的组合数和快速幂计算结果：
     • ans = C(n - 1, k) * m % MOD * (m - 1)^(n - k - 1) % MOD。

时间和空间复杂度

1. 时间复杂度：

   • 预计算阶乘和逆阶乘的时间是 O(MX)，其中 MX 是预计算的最大值（这里是 100000）。
   • 计算组合数和快速幂的时间是 O(1)（因为预计算已经完成）。
   • 因此，总的时间复杂度是 O(MX)，即 O(100000)。

2. 空间复杂度：

   • 需要存储阶乘和逆阶乘数组 fact 和 invFact，大小均为 MX。
   • 因此，额外的空间复杂度是 O(MX)，即 O(100000)。

总结

• 通过组合数学将问题分解为选择相邻相等的位置和分配块的值。
• 预计算阶乘和逆阶乘以高效计算组合数。
• 最终结果是组合数与块值分配方案的乘积。
• 时间和空间复杂度均为 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()
	return int(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:
        return 0
    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))