CF1808E3 Minibuses on Venus (hard version) 题解

题意：

给定 \(n,k,m\) ，定义序列 \(\{a_i\}_{i=1}^n\) 的值域为 \([0,k)\) ，即 \(\forall i, ~a_i \in[0,k)\) 。

称一个序列 \(a\) 是幸运的，当且仅当存在至少一个 \(p~(1\le p \le n)\) 使得 \[ \sum_{i=1}^na_i - a_p\equiv a_p \pmod{k} \] 求幸运序列的总方案数模 \(m\) 的结果。

输入格式：

一行，三个整数 \(n,k,m\) 。

输出格式：

一行，即答案。

数据范围：

\(1 \le n \le 10^{18},~1\le k \le 2000,~10^8 \le m \le 10^9\) 。

记 \(s = \sum_{i=1}^n a_i\) ，考虑问题的反面，即 \(s \not\equiv 2a_i \pmod{k}\) 的序列

考虑枚举 \(s \in [0,k)\) 。由于当 \(k\) 是奇数/偶数时，\(s \equiv 2x \pmod{k}\) 分别有唯一解/两个解，故分类讨论。

1. \(k\) 为奇数

此时 \(s \equiv 2x \pmod{k}\) 有唯一解（注意我们枚举了 \(s\) ，故现在 \(s\) 是已知量）

考虑容斥。设 \(f_i\) 表示钦定 \(i\) 个位置为 \(x\) 的方案数，\(g_i\) 为恰好 \(i\) 个位置为 \(x\) 的方案数，那么 \[ f_i = \sum_{j=i}^n\binom{j}{i}g_j \] 我们要求 \(g(0)\) 。考虑二项式反演，有 \[ g_i = \sum_{j=i}^n(-1)^{j-i}\binom{j}{i}f_j \] 代入可得 \[ g_0=\sum_{i=0}^n(-1)^i f_i \] 现在考虑计算 \(f_i\) 。可以发现，当 \(i=n\) 时，\(f_i = [nx \equiv s \pmod{k}]\) 。

否则，剩余的 \(n-i\) 个位置中有 \(n-i-1\) 个可以随便选，而最后一个要把和凑成 \(x\) ，故 \[ f_i = \binom{n}{i}k^{n-1-i} \] 那么代回去可以得到 \[ g_0 = \sum_{i=0}^{n-1}(-1)^i\binom{n}{i}k^{n-1-i} + (-1)^nf_n \] 对左侧考虑二项式定理，可以凑出 \[ g_0 = \frac{(k-1)^n - (-1)^n}{k} + (-1)^nf_n \] 注意到这个是建立在枚举 \(s\) 的基础上的。考虑重新将 \(f_n\) 记为 \(f_s(n)\) 。

对 \(s \in[0,k)\) 累加求和可以得到 \[ (k-1)^n-(-1)^n+(-1)^n \sum_{i=0}^{k-1} f_i(n) \] 再利用 \(x\) 的定义即可得到 \[ f_i(n) = [(n-2)i\equiv 0\pmod{k}] \] 这样的 \(i\) 显然有 \(\gcd(n - 2,\,k)\) 个。代入，并用总方案数减去它，得到该部分的方案数为 \[ k^n - (k-1)^n + (-1)^n - (-1)^n\gcd(n - 2, \, k) \]

2. \(k\) 为偶数

此时当 \(s\) 为奇数时 \(s \not\equiv 2a_i \pmod{k}\) 显然无解。

考虑统计 \(s\) 为偶数且不包含满足 \(2a_i\equiv 0\pmod{k}\) 的 \(a_i\) 的方案数。

不同于 \(k\) 为奇数的情况，此时不满足条件的数，即满足 \(s \equiv 2x \pmod{k}\) 的 \(x\) 有两个

不妨记为 \(x_1,x_2\) ，可以解出这两个值，分别为 \(x_1=\frac{s}{2},~x_2 = \frac{s + k}{2}\) 。

继续考虑容斥。设 \(f_i\) 表示钦定 \(i\) 个值为 \(x_1\) 或 \(x_2\) 的方案数，\(g_i\) 表示恰好 \(i\) 个值为 \(x_1\) 或 \(x_2\) 的方案数

我们要求的就是 \(g_0\) ，根据之前的推导，同理可得 \[ g_0 = \sum_{i=0}^n(-1)^if_i \] 考虑计算 \(f_i\) 。当 \(i \ne n\) 时 \(f_i = \binom{n}{i}2^ik^{n-1-i}\) ，否则记为 \(f_s(n)\) ，累加得 \[ \frac{(k-2)^n-(-2)^n}{2}+(-1)^n \sum_{0 \leq s=2x<k} f_s(n) \] 考虑计算 \(f_s(n)\) 。枚举其中有 \(j\) 个 \(x_1\) ，\(n-j\) 个 \(x_2\) ，那么 \[ f_s(n) = \sum_{j=0}^n\binom{n}{j}[jx_1 + (n-j)x_2\equiv 0 \pmod{k}] \] 代入 \(x_1,x_2\) ，化简可得 \[ f_s(n) = \sum_{j=0}^n\binom{n}{j}\left[\frac{k}{2}j\equiv\frac{s+k}{2}n-s\pmod{k}\right] \] 提示：这里中括号和之前的一样，也是艾弗森括号。

注意到 \(\frac{k}{2}j\) 模 \(k\) 的值只可能是 \(0\) 或 \(\frac{k}{2}\) ，而右侧是定值。

并且无论哪种情况均为组合数上指标的交换求和，其值均为 \(2^{n-1}\) 。

而 \(\frac{s+k}{2}n - s \equiv 0 \pmod{\frac{k}{2}}\) ，故 \[ \frac{s}{2}(n-2)\equiv 0 \pmod {\textstyle \frac{k}{2}} \] 这种情况下，满足条件的 \(s\) 有 \(\gcd(n-2,\,\frac{k}{2})\) 个。于是综上可得 \[ \sum_{0 \le s = 2x \le k}f_S(n)=2^{n-1} \operatorname{gcd}\left(n-2,\, \frac{k}{2}\right) \] 注意这是反面情况，那么整理可得 \[ \frac{k^n-(k-2)^n+(-2)^n}{2}-(-1)^n 2^{n-1} \operatorname{gcd}\left(n-2,\, \frac{k}{2}\right) \] 实现时要注意 \(2\) 可能没有模意义下的逆元，具体见代码。

时间复杂度 \(\mathcal{O}(\log n + \log k)\)

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
void up(int &x, int y) { x < y ? x = y : 0; }
void down(int &x, int y) { x > y ? x = y : 0; }
#define rep(i, a, b) for(int i = (a), i##END = (b); i <= i##END; i++)
#define Rep(i, a, b) for(int i = (a), i##END = (b); i >= i##END; i--)
#define N ((int)())

int n, k, mod;
int gcd(int a, int b) { return (!b) ? a : gcd(b, a % b); }
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 sgn(int x) { return (x & 1) ? mod - 1 : 1; } // (-1)^x
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> n >> k >> mod;
    if(n == 1 || k == 1) { return cout << "1\n", 0; }
    if(k & 1)
    {
        const int _1 = (qpow(k, n) + mod - qpow(k - 1, n)) % mod;
        const int _2 = (sgn(n) + mod - sgn(n) * gcd(n - 2, k) % mod) % mod;
        cout << (_1 + _2) % mod << '\n';
    }else
    {
        int t = k / 2;
        const int _1 = t * qpow(k, n - 1) % mod;
        const int _2 = (mod - (t - 1) * qpow(k - 2, n - 1) % mod) % mod;
        const int _3 = qpow(2, n - 1) * sgn(n) % mod;
        const int _4 = (mod - sgn(n) * qpow(2, n - 1) % mod * gcd(n - 2, t) % mod) % mod;
        cout << ((_1 + _2) % mod + (_3 + _4) % mod) % mod << '\n';
    }
    return 0;
}