2025-07-14:统计恰好有 K 个相等相邻元素的数组数目。用go语言,

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 = 3k = 1 时,数组可以分成 2 个块(因为 n - k = 2),如 [1, 1, 2](块 [1, 1][2])。

2. 组合数学

o 首先,我们需要选择 k 个位置作为相邻相等的位置。这相当于在 n - 1 个可能的位置中选择 k 个位置,使得这些位置是相邻相等的。这可以通过组合数 C(n - 1, k) 来计算。

o 然后,我们需要为这些“块”分配值。第一个块有 m 种选择(1m),后续的每个新块必须与前一个块的值不同,因此有 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 = 3m = 2k = 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 需要存储阶乘和逆阶乘数组 factinvFact,大小均为 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 语言和算法助力您的职业发展

·

原文链接:,转发请注明来源!