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