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;
}