CF1227F2 Wrong Answer on test 233 (Hard Version) 题解

题目链接：Wrong Answer on test 233 (Hard Version)

题意：

cxy 要参加一场考试，这场考试有 $n$ 道选择题，每个选择题有 $k$ 个选项。

但是 cxy 昨天玩某中文二字英文七字游戏玩过头了，导致她这次考试时十分困倦以至于居然填错了答题卡！

具体地，cxy 错误地将第 $i$ 道选择题的答案填到第 $(i \bmod n) + 1$ 道题去了。

~~不过她是怎么做到将第 $n$ 道题的答案填到第 $1$ 道去的？果然是玩【数据删除】玩的。~~

显然，此时 cxy 非常焦虑，但是她不记得自己每道题原来填了什么选项了

现给出每道题的答案 $h_i$ ，请帮助 cxy 计算有多少种可能的初始答案，使得错误填涂后比原来正确的多了。

输入格式：

第一行给出 $n,k~(1 \le n \le 2\times 10^5,\,1\le k \le 10^9)$ 。

第二行给出 $n$ 个数 $h_1,h_2,\cdots,h_n(1 \le h_i \le k)$ ，表示每道题的答案。

输出格式：

仅一行，输出方案数模 $998244353$ 的结果。

稍微魔改了一下题面，应该比原题更好理解了。

首先考虑 $n\le 2000$ 的简单版，这个直接 dp 即可。

设 $f_{i,j}$ 为填完前 $i$ 道题，填错答案后比原来多对了 $j(0\le j \le n)$ 个题的方案数。

若 $a_i = a_{(i \bmod n)+1}$ ，则 $f_{i,j} = f_{i-1,j}\times k$ 。
否则 $f_{i, j}=f_{i-1, j+1}+f_{i-1, j-1}+f_{i-1, j} \times(k-2)$ 。这里 $k-2$ 是因为不能填 $a_i$ 或 $a_{(i \bmod n)+1}$ 。

现在考虑 $n \le 2\times 10^5$ 。尽管这个dp式子并不好优化，它还是提供了一些有价值的性质。

不妨扩充 $j$ 的范围到 $[-n,n]$ ，可以发现此时递推式仍然和之前一样，于是 $f_{i,j}=f_{i,-j}$ 。

其实不用重新推一遍递推式，在转移过程中我们就可以发现，当前这个数贡献 $+1$ 或 $-1$ 方案数是一样的。

现在重新设 $f_i$ 为多对了 $i$ 道题的方案数，由对称性可知 $f_i = f_{-i}$ 。

那么我们的答案可以写作

$\mathrm{Ans} = \frac{k^n - f_0}{2}$

于是我们只需要求出 $f_0$ 就可以了。

设有 $m$ 个位置满足 $a_i \ne a_{(i \bmod n)+1}$ ，那么枚举对了多少题（等于错了多少题），有

$f_0 = k^{n - m}\sum_{i=0}^{\left\lfloor\frac{m}{2}\right\rfloor}\binom{m}{i}\binom{m-i}{i}(k-2)^{m - 2i}$

那么把 $f_0$ 算出来代到 $\mathrm{Ans}$ 里就好了。

时间复杂度 $\mathcal{O}(n)$

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
const int mod = 998244353;
void up(int &x, int y) { x < y ? x = y : 0; }
void down(int &x, int y) { x > y ? x = y : 0; }
void add(int &x, int y) { (x += y) >= mod ? x -= mod : 0; }
#define N ((int)(2e5 + 15))

int m, a[N], fac[N], infac[N], pw[N];
int qpow(int a, int b)
{
    int r = 1;
    for(a %= mod; b; b >>= 1, a = a * a % mod)
        if(b & 1) r = r * a % mod;
    return r;
}
int C(int n, int m)
{
    if(n < m || m < 0) return 0;
    return fac[n] * infac[m] % mod * infac[n - m] % mod;
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    int n, k; cin >> n >> k;
    if(k == 1) return cout << "0\n", 0;
    for(int i = 1; i <= n; i++) cin >> a[i];
    for(int i = 1; i <= n; i++) m += (a[i] != a[i % n + 1]);
    fac[0] = 1;
    for(int i = 1; i <= m; i++) fac[i] = fac[i - 1] * i % mod;
    infac[m] = qpow(fac[m], mod - 2);
    for(int i = m - 1; ~i; i--) infac[i] = infac[i + 1] * (i + 1) % mod;
    pw[0] = 1;
    for(int i = 1; i <= n; i++) pw[i] = pw[i - 1] * (k - 2) % mod;
    int res = qpow(k, n), w = qpow(k, n - m);
    for(int i = 0; i <= m / 2; i++) 
        add(res, mod - w * C(m, i) % mod * C(m - i, i) % mod * pw[m - 2 * i] % mod);
    cout << res * qpow(2, mod - 2) % mod << '\n';
    return 0;
}