晚饭吃什么？

晚饭吃什么？

题目描述

“晚饭吃什么？” “我都行。”

notPP 罗列了 $m$ 个吃饭的地点，在接下来的 $n$ 天中，notPP 觉得一直吃外卖会腻，于是他决定拿恰好 $k$ 天出来去外面吃饭。

但连续两天都在同一个地方吃饭容易腻，于是他希望“不能连续两天都去同一个地方吃饭（外卖除外）”的规则。

在这种情况下，请问 notPP 有多少种方案呢？因为答案可能很大，所以请将其对 $10^9+7$ 取模。

输入描述:

本题每个测试点有多组测试数据。

第一行输入一个正整数 $T$ $(1 \leq T \leq 10^5)$，表示有 $T$ 组测试样例。

接下来对于每组样例：

一行输入三个数字 $n,m,k$ $(1 \leq k \leq n \leq 10^6, 1 \leq m \leq 10^6)$，分别代表接下来 $n$ 天，有 $m$ 个吃饭的地方和分配 $k$ 天。

保证单个测试点的所有测试数据的 $n$ 之和不超过 $10^6$。

输出描述:

输出一行，表示分配方案的数量对 $10^9+7$ 取模得到的结果。

示例1

输入

2
3 2 3
3 2 2

输出

2
8

备注:

对于第一个样例，notPP 在接下来的 $3$ 天都去外面吃，故有两种情况：${1,2,1}$ 和 ${2,1,2}$。

对于第二个样例，这里有所有的可行情况（其中 $0$ 表示外卖）：

$\{1,2,0\},\{2,1,0\},\{1,0,1\},\{1,0,2\},\{2,0,1\},\{2,0,2\},\{0,1,2\},\{0,2,1\}" title="\{1,2,0\},\{2,1,0\},\{1,0,1\},\{1,0,2\},\{2,0,1\},\{2,0,2\},\{0,1,2\},\{0,2,1\}$

 

解题思路

  题目等价于，从 $n$ 个空白格子中选出 $k$ 个格子进行染色，其中相邻的染色格子要求颜色不同，一共有 $m$ 种不同的颜色。假设有一段连续的染色格子，长度为 $c$，由乘法原理知道这段格子的染色方案数为 $m \cdot (m-1)^{c-1}$。因此如果有 $s \, (1 \leq s \leq k)$ 段这样不相连的连续染色格子，那么染色的方案数就是 $\prod\limits_{i=1}^{s}{m \cdot (m-1)^{c_i-1}} = m^s \cdot (m-1)^{k-s}$。

  剩下的问题是在 $n$ 个格子中可以划分出多少种 $s$ 段不相连的连续染色格子。反过来思考，我们先把 $k$ 个格子划分成 $s$ 段，方案数是 $C_{k-1}^{s-1}$（相当于从 $k$ 个格子选出 $s$ 个作为每一段开头，由于第一个格子一定是第一段的开头，因此只需要从剩下 $k-1$ 个格子选择 $s-1$ 个作为开头），然后再把这 $s$ 段分别插入到由 $n-k$ 个空白格子构成的 $n-k+1$ 个间隙中（每个间隙只能插入一段），方案数是 $C_{n-k+1}^{s}$，因此总的方案数就是 $C_{k-1}^{s-1} \cdot C_{n-k+1}^{s}$。

  所以当有 $s$ 段不相连的连续染色格子时，总的染色方案数是 $m^s \cdot (m-1)^{k-s} \cdot C_{k-1}^{s-1} \cdot C_{n-k+1}^{s}$。枚举所有可能的 $s$，总的答案就是 $\sum\limits_{s=1}^{k}{m^s \cdot (m-1)^{k-s} \cdot C_{k-1}^{s-1} \cdot C_{n-k+1}^{s}}$。

  AC 代码如下，时间复杂度为 $O(n)$：

#include <bits/stdc++.h>
using namespace std;

typedef long long LL;

const int N = 1e6 + 5, mod = 1e9 + 7;

int fact[N], inv[N];
int p1[N], p2[N];

int C(int a, int b) {
    if (a < b) return 0;
    return 1ll * fact[a] * inv[b] % mod * inv[a - b] % mod;
}

void solve() {
    int n, m, k;
    cin >> n >> m >> k;
    fact[0] = inv[0] = p1[0] = p2[0] = 1;
    for (int i = 1; i <= n; i++) {
        fact[i] = 1ll * fact[i - 1] * i % mod;
        p1[i] = 1ll * p1[i - 1] * m % mod;
        p2[i] = p2[i - 1] * (m - 1ll) % mod;
    }
    inv[1] = 1;
    for (int i = 2; i <= n; i++) {
        inv[i] = LL(mod - mod / i) * inv[mod % i] % mod;
    }
    for (int i = 2; i <= n; i++) {
        inv[i] = 1ll * inv[i] * inv[i - 1] % mod;
    }
    int ret = 0;
    for (int i = 1; i <= k; i++) {
        ret = (ret + 1ll * p1[i] * p2[k - i] % mod * C(n - k + 1, i) % mod * C(k - 1, i - 1)) % mod;
    }
    cout << ret << '\n';
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int t;
    cin >> t;
    while (t--) {
        solve();
    }
    
    return 0;
}

 

参考资料

  “华为杯”2025 年广东工业大学 ACM 程序设计竞赛题解：https://ac.nowcoder.com/discuss/1481277