题目描述

给你三个整数 n，m，k。长度为 n 的 好数组 arr 定义如下：

• arr 中每个元素都在 闭 区间 [1, m] 中。
• 恰好 有 k 个下标 i （其中 1 <= i < n）满足 arr[i - 1] == arr[i]。

请你返回可以构造出的 好数组 数目。

由于答案可能会很大，请你将它对 10^9 + 7 取余 后返回。

样例

输入：n = 3, m = 2, k = 1

输出：4

解释：

总共有 4 个好数组，分别是 [1, 1, 2]，[1, 2, 2]，[2, 1, 1] 和 [2, 2, 1]。
所以答案为 4。

输入：n = 4, m = 2, k = 2

输出：6

解释：

好数组包括 [1, 1, 1, 2]，[1, 1, 2, 2]，[1, 2, 2, 2]
    [2, 1, 1, 1]，[2, 2, 1, 1] 和 [2, 2, 2, 1]。
所以答案为 6 。

输入：n = 5, m = 2, k = 0

输出：2

解释：

好数组包括 [1, 2, 1, 2, 1] 和 [2, 1, 2, 1, 2]。
所以答案为 2。

限制

• 1 <= n <= 10^5
• 1 <= m <= 10^5
• 0 <= k <= n - 1

算法

(组合数学) $O(n)$

1. 第一个位置有 $m$ 种选择。
2. 剩余 $n - 1$ 个位置中，选出 $k$ 个位置满足 $arr(i) = arr(i - 1)$，这对于选中的位置来说，只会有一种选择。
3. 其余 $n - 1 - k$ 个位置中，每个位置需要保证 $arr(i) \neq arr(i - 1)$，故有 $m - 1$ 种选择。
4. 综上，答案就是 $m \cdot C_{n - 1}^{k} \cdot (m - 1)^{n - 1 - k}$。
5. 可以在线性时间预处理 $1$ 到 $N$ 的阶乘和阶乘的对 $mod$ 的乘法逆元。

时间复杂度

• 预处理的时间复杂度为 $O(n)$。
• 快速幂计算答案的时间复杂度为 $O(\log n)$。
• 故总时间复杂度为 $O(n)$。

空间复杂度

• 需要 $O(n)$ 的额外空间存储阶乘和阶乘的逆元。

C++ 代码

#define LL long long

const int mod = 1000000007;
const int N = 100010;

int fac[N], inv_fac[N];

int power(int x, int y) {
    int res = 1, p = x;
    for (; y; y >>= 1) {
        if (y & 1) res = (LL)(res) * p % mod;
        p = (LL)(p) * p % mod;
    }

    return res;
}

auto init = []{
    fac[0] = 1;
    for (int i = 1; i < N; i++)
        fac[i] = (LL)(fac[i - 1]) * i % mod;

    inv_fac[N - 1] = power(fac[N - 1], mod - 2);
    for (int i = N - 2; i >= 0; i--)
        inv_fac[i] = (LL)(inv_fac[i + 1]) * (i + 1) % mod;

    return 0;
}();

class Solution {
private:
    inline int C(int x, int y) {
        return (LL)(fac[x]) * inv_fac[y] % mod * inv_fac[x - y] % mod;
    }

public:
    int countGoodArrays(int n, int m, int k) {
        return (LL)(m) * C(n - 1, k) % mod * power(m - 1, n - 1 - k) % mod;
    }
};